ui/tabs.py

import gradio as gr
import threading
import time
import os
import traceback
from services.audio_service import AudioService
from services.chat_service import ChatService
from services.image_service import ImageService
from services.streaming_voice_service import StreamingVoiceService
from core.rag_system import EnhancedRAGSystem
from core.tts_service import EnhancedTTSService
from core.wikipedia_processor import WikipediaProcessor
from ui.components import create_audio_components, create_chat_components,create_streaming_voice_components

def create_all_tabs(audio_service: AudioService, chat_service: ChatService, 
                   image_service: ImageService, rag_system: EnhancedRAGSystem,
                   tts_service: EnhancedTTSService, wikipedia_processor: WikipediaProcessor,
                   streaming_voice_service: StreamingVoiceService):
    
    with gr.Tab("🎙️ Streaming Voice (VAD)"):
        create_streaming_voice_tab(streaming_voice_service)
    
    with gr.Tab("🎙️ Audio"):
        create_audio_tab(audio_service)
    
    with gr.Tab("💬 Chat"):
        create_chat_tab(chat_service)
    
    with gr.Tab("🖼️ Image"):
        create_image_tab(image_service)
    
    with gr.Tab("📚 RAG Wikipedia"):
        create_rag_tab(rag_system, wikipedia_processor)
    
    with gr.Tab("🔊 Text-to-Speech"):
        create_tts_tab(tts_service)
    
    with gr.Tab("🌐 Language Info"):  # NEW TAB
        create_language_info_tab(rag_system.multilingual_manager)
def create_rag_tab(rag_system: EnhancedRAGSystem, wikipedia_processor: WikipediaProcessor):
    """Tạo tab RAG với debug chi tiết"""
    
    # Initialize systems if not provided
    if rag_system is None:
        rag_system = EnhancedRAGSystem()
    if wikipedia_processor is None:
        wikipedia_processor = WikipediaProcessor()
    
    with gr.Blocks() as rag_tab:
        gr.Markdown("## 📚 Upload Dữ Liệu Wikipedia")
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📤 Upload Dữ Liệu")
                file_upload = gr.File(
                    label="Tải lên file (TXT, CSV, JSON)",
                    file_types=['.txt', '.csv', '.json'],
                    file_count="single"
                )
                upload_btn = gr.Button("📤 Upload Data", variant="primary")
                upload_status = gr.Textbox(
                    label="Trạng thái Upload",
                    interactive=False,
                    lines=5
                )
                
                gr.Markdown("### 📊 Thống kê Database")
                stats_btn = gr.Button("📊 Database Stats", variant="secondary")
                stats_display = gr.Textbox(
                    label="Thống kê", 
                    interactive=False,
                    lines=6
                )
                
            with gr.Column(scale=2):
                gr.Markdown("### 🔍 Tìm kiếm & Kiểm tra")
                search_query = gr.Textbox(
                    label="Tìm kiếm trong database",
                    placeholder="Nhập từ khóa để kiểm tra dữ liệu..."
                )
                search_btn = gr.Button("🔍 Tìm kiếm", variant="secondary")
                rag_results = gr.JSON(
                    label="Kết quả tìm kiếm",
                    show_label=True
                )
        
        def upload_wikipedia_file(file):
            """Xử lý upload file với debug đầy đủ"""
            if file is None:
                return "❌ Vui lòng chọn file để upload"
            
            try:
                print(f"🔄 Bắt đầu upload file: {file.name}")
                
                # Kiểm tra file tồn tại
                if not os.path.exists(file.name):
                    return f"❌ File không tồn tại: {file.name}"
                
                # Xử lý file
                documents = wikipedia_processor.process_uploaded_file(file.name)
                
                if not documents:
                    return "❌ Không thể trích xuất dữ liệu từ file. File có thể trống hoặc định dạng không đúng."
                
                print(f"✅ Đã xử lý {len(documents)} documents")
                
                # Tạo metadata
                metadatas = []
                for i, doc in enumerate(documents):
                    metadata = {
                        "source": "uploaded_file",
                        "type": "knowledge", 
                        "file_name": os.path.basename(file.name),
                        "language": "vi",
                        "doc_id": i,
                        "length": len(doc)
                    }
                    metadatas.append(metadata)
                
                # Thêm vào RAG system
                old_stats = rag_system.get_collection_stats()
                old_count = old_stats['total_documents']
                
                rag_system.add_documents(documents, metadatas)
                
                # Lấy thống kê mới
                new_stats = rag_system.get_collection_stats()
                new_count = new_stats['total_documents']
                
                success_msg = f"""
    ✅ UPLOAD THÀNH CÔNG!
    📁 File: {os.path.basename(file.name)}
    📄 Documents xử lý: {len(documents)}
    📊 Documents thêm vào: {new_count - old_count}
    🏷️ Tổng documents: {new_count}
    🔤 Embeddings: {new_stats['embedding_count']}
    🌐 Ngôn ngữ: {new_stats['language_distribution']}
    💡 Bạn có thể tìm kiếm ngay để kiểm tra dữ liệu!
    """
                return success_msg
                
            except Exception as e:
                error_msg = f"❌ LỖI UPLOAD: {str(e)}"
                print(f"UPLOAD ERROR: {traceback.format_exc()}")
                return error_msg
        
        def get_rag_stats():
            """Lấy thống kê chi tiết"""
            try:
                stats = rag_system.get_collection_stats()
                return f"""
    📊 THỐNG KÊ RAG DATABASE:
    • 📄 Tổng documents: {stats['total_documents']}
    • 🔤 Số embeddings: {stats['embedding_count']}
    • 📐 Dimension: {stats['embedding_dimension']}
    • 🌐 Phân bố ngôn ngữ: {stats['language_distribution']}
    • ✅ Trạng thái: {stats['status']}
    • 🏷️ Tên: {stats['name']}
    💡 Embeddings: {'Có' if stats['has_embeddings'] else 'Không'}
    """
            except Exception as e:
                return f"❌ Lỗi lấy thống kê: {str(e)}"
        
        def search_rag_database(query):
            """Tìm kiếm để kiểm tra dữ liệu"""
            if not query.strip():
                return [{"message": "Nhập từ khóa để tìm kiếm"}]
            
            try:
                results = rag_system.semantic_search(query, top_k=3)
                
                if not results:
                    return [{"message": "Không tìm thấy kết quả nào", "query": query}]
                
                return results
                
            except Exception as e:
                return [{"error": f"Lỗi tìm kiếm: {str(e)}"}]
        
        # Event handlers
        upload_btn.click(upload_wikipedia_file, inputs=[file_upload], outputs=[upload_status])
        stats_btn.click(get_rag_stats, inputs=[], outputs=[stats_display])
        search_btn.click(search_rag_database, inputs=[search_query], outputs=[rag_results])
    
    return rag_tab
def create_audio_tab(audio_service: AudioService):
    gr.Markdown("## Nói chuyện với AI (Đa ngôn ngữ)")
    audio_input, transcription_output, response_output, tts_audio_output, process_button = create_audio_components()
    
    # NEW: Language display
    language_display = gr.Textbox(
        label="🌐 Ngôn ngữ phát hiện",
        interactive=False,
        placeholder="Ngôn ngữ sẽ hiển thị ở đây..."
    )
    
    process_button.click(
        audio_service.transcribe_audio,
        inputs=audio_input,
        outputs=[transcription_output, response_output, tts_audio_output, language_display]  # UPDATED
    )
def create_streaming_voice_tab(streaming_service: StreamingVoiceService):
    """Tạo tab streaming voice với VAD optimized - FIXED VERSION"""
    
    with gr.Blocks() as streaming_tab:
        gr.Markdown("## 🎤 Trò chuyện giọng nói thời gian thực - Tối ưu hóa")
        
        # FIXED: Thêm state để lưu kết quả VAD
        vad_result_state = gr.State(value=None)
        
        with gr.Row():
            with gr.Column(scale=1):
                # Controls
                with gr.Row():
                    start_btn = gr.Button("🎙️ Bắt đầu VAD", variant="primary")
                    stop_btn = gr.Button("🛑 Dừng VAD", variant="secondary")
                
                gr.Markdown("### Chế độ tự động (VAD)")
                gr.Markdown("Hệ thống tự động nhận diện khi bạn bắt đầu nói")
                
                with gr.Row():
                    vad_status = gr.Textbox(
                        label="Trạng thái VAD",
                        value="Chưa bắt đầu",
                        interactive=False
                    )
                
                gr.Markdown("### Chế độ thủ công")
                microphone = gr.Microphone(
                    label="🎤 Nhấn để nói thủ công",
                    type="numpy",
                    streaming=True
                )
                
                with gr.Accordion("📊 Performance Metrics", open=False):
                    latency_display = gr.JSON(
                        label="Latency Statistics",
                        value={}
                    )
                    refresh_latency_btn = gr.Button("🔄 Refresh Metrics", size="sm")
                
                clear_btn = gr.Button("🗑️ Xóa hội thoại")
                
                # State info
                state_info = gr.Textbox(
                    label="Thông tin hệ thống",
                    value="Khởi tạo...",
                    lines=3,
                    interactive=False
                )
            
            with gr.Column(scale=2):
                transcription_box = gr.Textbox(
                    label="📝 Bạn vừa nói",
                    lines=2,
                    interactive=False
                )
                
                response_box = gr.Textbox(
                    label="🤖 Phản hồi AI",
                    lines=3,
                    interactive=False
                )
                
                audio_output = gr.Audio(
                    label="🔊 Giọng nói AI",
                    interactive=False,
                    autoplay=True
                )
        
        # State variables
        is_vad_active = gr.State(value=False)
        
        def vad_callback(result):
            """Callback khi VAD phát hiện speech - FIXED VERSION"""
            print(f"🎯 Nhận kết quả từ VAD: {result['transcription']}")
            vad_result_state.value = result
        
        def start_vad():
            """Bắt đầu VAD - FIXED VERSION"""
            success = streaming_service.start_listening(vad_callback)
            if success:
                is_vad_active.value = True
                status = "✅ VAD đang chạy - Hãy nói gì đó!"
                state = streaming_service.get_conversation_state()
                state_text = f"VAD: Đang hoạt động\nQueue: {state['queue_size']}\nThreads: {state['worker_threads']}"
            else:
                status = "❌ Không thể khởi động VAD"
                state_text = "Lỗi khởi động"
            
            return status, state_text
        
        def stop_vad():
            """Dừng VAD"""
            streaming_service.stop_listening()
            is_vad_active.value = False
            state = streaming_service.get_conversation_state()
            state_text = f"VAD: Đã dừng\nHistory: {state['history_length']} messages"
            return "🛑 VAD đã dừng", state_text
        
        def process_microphone(audio_data):
            """Xử lý microphone input"""
            if audio_data is None:
                return "Chưa có âm thanh", "Hãy nói gì đó...", None, "VAD: Đang chạy" if is_vad_active.value else "VAD: Dừng"
            
            try:
                result = streaming_service.process_streaming_audio(audio_data)
                state = streaming_service.get_conversation_state()
                state_text = f"Manual mode\nHistory: {state['history_length']} messages"
                return result['transcription'], result['response'], result['tts_audio'], state_text
            except Exception as e:
                return f"Lỗi: {e}", "Xin lỗi, có lỗi xảy ra", None, "Lỗi xử lý"
        
        def check_vad_results():
            """Kiểm tra kết quả VAD - FIXED: Cập nhật UI khi có kết quả mới"""
            if vad_result_state.value is not None:
                result = vad_result_state.value
                vad_result_state.value = None  # Reset sau khi dùng
                
                state = streaming_service.get_conversation_state()
                state_text = f"VAD mode\nHistory: {state['history_length']} messages\nQueue: {state['queue_size']}"
                
                return result['transcription'], result['response'], result['tts_audio'], state_text
            return gr.skip(), gr.skip(), gr.skip(), gr.skip()
        
        def clear_chat():
            """Xóa hội thoại"""
            streaming_service.clear_conversation()
            state = streaming_service.get_conversation_state()
            state_text = f"Đã xóa hội thoại\nHistory: {state['history_length']} messages"
            return "", "", None, state_text
        
        def refresh_latency():
            """Làm mới latency metrics"""
            stats = streaming_service.get_latency_stats()
            return stats
        
        def update_state_info():
            """Cập nhật thông tin trạng thái"""
            state = streaming_service.get_conversation_state()
            
            formatted_state = f"VAD: {'Đang chạy' if state['is_listening'] else 'Dừng'}\n"
            formatted_state += f"Queue: {state['queue_size']}\n"
            formatted_state += f"History: {state['history_length']} messages\n"
            formatted_state += f"Threads: {state['worker_threads']}\n"
            formatted_state += f"Last: {state['last_update']}"
            
            latency_info = streaming_service.get_latency_stats()
            
            return formatted_state, latency_info

        # Event handlers
        start_btn.click(start_vad, outputs=[vad_status, state_info])
        stop_btn.click(stop_vad, outputs=[vad_status, state_info])
        
        microphone.stream(
            process_microphone,
            inputs=[microphone],
            outputs=[transcription_box, response_box, audio_output, state_info]
        )
        
        clear_btn.click(
            clear_chat,
            outputs=[transcription_box, response_box, audio_output, state_info]
        )

        refresh_latency_btn.click(
            refresh_latency,
            outputs=[latency_display]
        )
        
        # FIXED: Timer để kiểm tra kết quả VAD
        gr.Timer(1.0).tick(
            fn=check_vad_results,
            outputs=[transcription_box, response_box, audio_output, state_info]
        )
        
        # Timer để cập nhật system info
        gr.Timer(3.0).tick(
            fn=update_state_info,
            outputs=[state_info, latency_display]
        )
    
    return streaming_tab
def create_image_tab(image_service: ImageService):
    """Tạo tab phân tích hình ảnh với OCR và LLM"""
    
    with gr.Blocks() as image_tab:
        gr.Markdown("## 🖼️ Phân tích hình ảnh & Trích xuất văn bản")
        gr.Markdown("""
        ### 🔍 Chức năng:
        - **OCR đa ngôn ngữ**: Trích xuất văn bản từ ảnh (Tiếng Việt, Anh, Nhật, Hàn, Trung, ...)
        - **Phân tích AI**: Sử dụng LLM để phân tích nội dung và ngữ cảnh
        - **Hỗ trợ nhiều định dạng**: Tài liệu, ảnh chụp, meme, screenshot
        
        ### 📝 Hướng dẫn:
        1. Tải lên hình ảnh có chứa văn bản
        2. (Tùy chọn) Mô tả hình ảnh để AI phân tích chính xác hơn
        3. Nhấn "Phân tích hình ảnh" để xem kết quả
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                # Image input
                image_input = gr.Image(
                    type="numpy", 
                    label="🖼️ Tải lên hình ảnh",
                    height=300
                )
                
                # User description
                image_description = gr.Textbox(
                    label="📝 Mô tả hình ảnh (tùy chọn)",
                    placeholder="Ví dụ: Đây là hóa đơn mua hàng, ảnh chụp menu nhà hàng, văn bản tiếng Việt...",
                    lines=3
                )
                
                # Analyze button
                analyze_button = gr.Button(
                    "🔍 Phân tích hình ảnh", 
                    variant="primary",
                    size="lg"
                )
                
                # Clear button
                clear_button = gr.Button("🗑️ Xóa", variant="secondary")
            
            with gr.Column(scale=2):
                # Results output
                image_output = gr.Textbox(
                    label="📊 Kết quả phân tích",
                    lines=15,
                    max_lines=20,
                    show_copy_button=True
                )
        
        def analyze_image(image, description):
            """Xử lý phân tích ảnh"""
            if image is None:
                return "❌ Vui lòng tải lên hình ảnh trước khi phân tích."
            
            return image_service.analyze_image_with_description(image, description)
        
        def clear_all():
            """Xóa tất cả input và output"""
            return None, "", ""
        
        # Event handlers
        analyze_button.click(
            analyze_image,
            inputs=[image_input, image_description],
            outputs=[image_output]
        )
        
        clear_button.click(
            clear_all,
            outputs=[image_input, image_description, image_output]
        )
    
    return image_tab
def create_chat_tab(chat_service: ChatService):
    gr.Markdown("## Trò chuyện với AI Assistant (Đa ngôn ngữ)")
    
    chatbot, state, user_input, send_button, clear_button, chat_tts_output = create_chat_components()
    
    # Language display
    chat_language_display = gr.Textbox(
        label="🌐 Ngôn ngữ phát hiện",
        interactive=False,
        placeholder="Ngôn ngữ sẽ hiển thị ở đây..."
    )
    
    # SỬA: Đảm bảo số lượng inputs/outputs khớp
    send_button.click(
        fn=chat_service.respond,
        inputs=[user_input, state],
        outputs=[user_input, chatbot, state, chat_tts_output, chat_language_display]
    )
    
    clear_button.click(
        fn=chat_service.clear_chat_history,
        inputs=[state],
        outputs=[chatbot, state]
    )
    
    # Thêm enter để gửi tin nhắn
    user_input.submit(
        fn=chat_service.respond,
        inputs=[user_input, state],
        outputs=[user_input, chatbot, state, chat_tts_output, chat_language_display]
    )

def create_language_info_tab(multilingual_manager):
    """Tab hiển thị thông tin về hệ thống đa ngôn ngữ"""
    gr.Markdown("## 🌐 Thông tin Hệ thống Đa ngôn ngữ")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🔧 Cấu hình Model")
            
            vietnamese_info = multilingual_manager.get_language_info('vi')
            multilingual_info = multilingual_manager.get_language_info('en')
            
            # SỬA LỖI Ở ĐÂY: Sử dụng key đúng từ dictionary mới
            gr.Markdown(f"""
            **Tiếng Việt:**
            - Embedding Model: `{vietnamese_info['embedding_model']}`
            - LLM Model: `{vietnamese_info['llm_model']}`
            - Trạng thái: {vietnamese_info['embedding_status']}
            
            **Đa ngôn ngữ:**
            - Embedding Model: `{multilingual_info['embedding_model']}`
            - LLM Model: `{multilingual_info['llm_model']}`
            - Trạng thái: {multilingual_info['embedding_status']}
            """)
        
        with gr.Column():
            gr.Markdown("### 🎯 Ngôn ngữ được hỗ trợ")
            
            supported_languages = """
            - 🇻🇳 **Tiếng Việt**: Sử dụng model chuyên biệt
            - 🇺🇸 **English**: Sử dụng model đa ngôn ngữ
            - 🇫🇷 **French**: Sử dụng model đa ngôn ngữ
            - 🇪🇸 **Spanish**: Sử dụng model đa ngôn ngữ
            - 🇩🇪 **German**: Sử dụng model đa ngôn ngữ
            - 🇯🇵 **Japanese**: Sử dụng model đa ngôn ngữ
            - 🇰🇷 **Korean**: Sử dụng model đa ngôn ngữ
            - 🇨🇳 **Chinese**: Sử dụng model đa ngôn ngữ
            """
            gr.Markdown(supported_languages)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🔍 Kiểm tra Ngôn ngữ")
            test_text = gr.Textbox(
                label="Nhập văn bản để kiểm tra ngôn ngữ",
                placeholder="Nhập văn bản bằng bất kỳ ngôn ngữ nào..."
            )
            test_button = gr.Button("🔍 Kiểm tra", variant="primary")
            
            test_result = gr.JSON(label="Kết quả phát hiện ngôn ngữ")
        
        test_button.click(
            lambda text: {
                'detected_language': multilingual_manager.detect_language(text),
                'language_info': multilingual_manager.get_language_info(multilingual_manager.detect_language(text)),
                'embedding_model': multilingual_manager.get_embedding_model(multilingual_manager.detect_language(text)) is not None,
                'llm_model': multilingual_manager.get_llm_model(multilingual_manager.detect_language(text))
            },
            inputs=[test_text],
            outputs=[test_result]
        )
def create_tts_tab(tts_service: EnhancedTTSService):
        gr.Markdown("## 🎵 Chuyển văn bản thành giọng nói nâng cao")
        gr.Markdown("Nhập văn bản và chọn ngôn ngữ để chuyển thành giọng nói")
        
        with gr.Group():
            with gr.Row():
                tts_text_input = gr.Textbox(
                    label="Văn bản cần chuyển thành giọng nói",
                    lines=4,
                    placeholder="Nhập văn bản tại đây..."
                )
            with gr.Row():
                tts_language = gr.Dropdown(
                    choices=["vi", "en", "fr", "es", "de", "ja", "ko", "zh"],
                    value="vi",
                    label="Ngôn ngữ"
                )
                tts_provider = gr.Dropdown(
                    choices=["auto", "gtts", "edgetts"],
                    value="auto",
                    label="Nhà cung cấp TTS"
                )
            with gr.Row():
                tts_output_audio = gr.Audio(
                    label="Kết quả giọng nói",
                    interactive=False
                )
            tts_button = gr.Button("🔊 Chuyển thành giọng nói", variant="primary")
        
        def text_to_speech_standalone(text, language, tts_provider):
            if not text:
                return None
            
            try:
                tts_audio_bytes = tts_service.text_to_speech(text, language, tts_provider)
                if tts_audio_bytes:
                    temp_audio_file = tts_service.save_audio_to_file(tts_audio_bytes)
                    return temp_audio_file
            except Exception as e:
                print(f"❌ Lỗi TTS: {e}")
            
            return None
        
        tts_button.click(
            text_to_speech_standalone,
            inputs=[tts_text_input, tts_language, tts_provider],
            outputs=[tts_output_audio]
        )