app.py

"""
Chainlit UI for AI Agent with Login/Registration - FIXED VERSION
"""

import chainlit as cl
import httpx  # Thay thế requests bằng httpx cho async support
from typing import Optional, Dict
import json
import os
import asyncio

# API Configuration
API_BASE_URL = os.getenv("API_BASE_URL")

# Timeout settings (giây)
API_TIMEOUT = 300.0  # 5 phút cho chat API (có thể điều chỉnh)
AUTH_TIMEOUT = 150.0  # 30 giây cho login/register

# ============================================
# API HELPER FUNCTIONS - CONVERTED TO ASYNC
# ============================================

async def check_user_exists(user_id: str) -> bool:
    """Check if user exists in the system"""
    try:
        async with httpx.AsyncClient(timeout=AUTH_TIMEOUT) as client:
            response = await client.get(f"{API_BASE_URL}/user-exists/{user_id}")
            if response.status_code == 200:
                return response.json().get("exists", False)
            return False
    except httpx.TimeoutException:
        print(f"Timeout checking user: {user_id}")
        return False
    except Exception as e:
        print(f"Error checking user: {e}")
        return False


async def register_user(user_id: str) -> Dict:
    """Register a new user"""
    try:
        async with httpx.AsyncClient(timeout=AUTH_TIMEOUT) as client:
            response = await client.post(
                f"{API_BASE_URL}/register",
                json={"user_id": user_id}
            )
            if response.status_code == 201:
                return {"success": True, "message": response.json().get("message", "Đăng ký thành công!")}
            else:
                error_msg = response.json().get("detail", "Đăng ký thất bại")
                return {"success": False, "message": error_msg}
    except httpx.TimeoutException:
        return {"success": False, "message": "Timeout: Server không phản hồi"}
    except httpx.ConnectError:
        return {"success": False, "message": "Không thể kết nối đến server"}
    except Exception as e:
        return {"success": False, "message": f"Lỗi: {str(e)}"}


async def send_chat_message(query: str, user_id: str) -> Dict:
    """Send chat message to API with streaming support"""
    try:
        # Sử dụng timeout dài hơn cho chat vì có thể xử lý lâu
        async with httpx.AsyncClient(timeout=API_TIMEOUT) as client:
            response = await client.post(
                f"{API_BASE_URL}/chat",
                json={"query": query, "user_id": user_id}
            )
            
            if response.status_code == 200:
                data = response.json()
                return {
                    "success": True,
                    "answer": data.get("answer", ""),
                    "references": data.get("references", "Không có tài liệu tham khảo")
                }
            else:
                error_msg = response.json().get("detail", "Lỗi xử lý câu hỏi")
                return {"success": False, "message": error_msg}
                
    except httpx.TimeoutException:
        return {
            "success": False, 
            "message": f"⏱️ Server đang xử lý quá lâu (>{API_TIMEOUT}s). Vui lòng thử lại hoặc đặt câu hỏi đơn giản hơn."
        }
    except httpx.ConnectError:
        return {"success": False, "message": "❌ Không thể kết nối đến server. Vui lòng kiểm tra API_BASE_URL."}
    except Exception as e:
        return {"success": False, "message": f"❌ Lỗi: {str(e)}"}


# ============================================
# CHAINLIT EVENT HANDLERS
# ============================================

@cl.on_chat_start
async def start():
    """Initialize chat session"""
    # Check if user is already logged in
    user_id = cl.user_session.get("user_id")
    
    if not user_id:
        await show_login_screen()
    else:
        await show_chat_interface(user_id)


async def show_login_screen():
    """Display login/registration screen"""
    welcome_msg = """
### 🎼 Chào mừng đến với Kumiko v2! 📯

**Đăng nhập** hoặc **đăng ký** để bắt đầu trò chuyện với mình nhé!

---

#### 📝 Hướng dẫn:
- **Đăng nhập**: Gõ `/login <user_id>` (ví dụ: `/login john`)
- **Đăng ký**: Gõ `/register <user_id>` (ví dụ: `/register john`)
- User ID phải có ít nhất 3 ký tự

---
**Ví dụ:**
- `/login alice`
- `/register bob123`
"""
    
    await cl.Message(content=welcome_msg, author="Kumiko").send()
    cl.user_session.set("awaiting_auth", True)


async def show_chat_interface(user_id: str):
    """Display chat interface after successful login"""
    welcome_msg = f"""#### ✅ Yayy, bạn đã đăng nhập thành công!

**User ID của bạn là:** `{user_id}`

---

Bạn có thể bắt đầu đặt câu hỏi. Mình sẽ trả lời và cung cấp tài liệu tham khảo nếu có!

💡 **Mẹo:** Sau mỗi câu trả lời, bạn có thể nhấn nút "𝄢 Tài liệu tham khảo" để check lại thông tin quan trọng!

---
Để đăng xuất, gõ `/logout`
"""
    
    await cl.Message(content=welcome_msg, author="Kumiko").send()
    cl.user_session.set("user_id", user_id)
    cl.user_session.set("awaiting_auth", False)


@cl.on_message
async def main(message: cl.Message):
    """Handle incoming messages"""
    user_id = cl.user_session.get("user_id")
    awaiting_auth = cl.user_session.get("awaiting_auth", True)
    
    content = message.content.strip()
    
    # Handle authentication commands
    if content.startswith("/"):
        await handle_command(content)
        return
    
    # Check if user is logged in
    if not user_id or awaiting_auth:
        await cl.Message(
            content="⚠️ Vui lòng đăng nhập trước khi chat.\n\nGõ `/login <user_id>` hoặc `/register <user_id>`",
            author="Kumiko"
        ).send()
        return
    
    # Process chat message
    await process_chat_message(content, user_id)


async def handle_command(command: str):
    """Handle special commands"""
    parts = command.split(maxsplit=1)
    cmd = parts[0].lower()
    
    # Login command
    if cmd == "/login":
        if len(parts) < 2:
            await cl.Message(content="❌ Sử dụng: `/login <user_id>`", author="Kumiko").send()
            return
        
        user_id = parts[1].strip()
        
        if len(user_id) < 3:
            await cl.Message(content="❌ User ID phải có ít nhất 3 ký tự!", author="Kumiko").send()
            return
        
        # Show loading message
        loading_msg = cl.Message(content="🔍 Đang kiểm tra tài khoản...", author="Kumiko")
        await loading_msg.send()
        
        # Check if user exists (now async)
        exists = await check_user_exists(user_id)
        
        await loading_msg.remove()
        
        if exists:
            await show_chat_interface(user_id)
        else:
            await cl.Message(
                content=f"❌ User ID `{user_id}` không tồn tại!\n\nVui lòng đăng ký bằng: `/register {user_id}`",
                author="Kumiko"
            ).send()
    
    # Register command
    elif cmd == "/register":
        if len(parts) < 2:
            await cl.Message(content="❌ Sử dụng: `/register <user_id>`", author="Kumiko").send()
            return
        
        user_id = parts[1].strip()
        
        if len(user_id) < 3:
            await cl.Message(content="❌ User ID phải có ít nhất 3 ký tự!", author="Kumiko").send()
            return
        
        # Show loading message
        loading_msg = cl.Message(content="📝 Đang đăng ký tài khoản...", author="Kumiko")
        await loading_msg.send()
        
        # Register user (now async)
        result = await register_user(user_id)
        
        await loading_msg.remove()
        
        if result["success"]:
            await cl.Message(content=f"✅ {result['message']}", author="Kumiko").send()
            await show_chat_interface(user_id)
        else:
            await cl.Message(content=f"❌ {result['message']}", author="Kumiko").send()
    
    # Logout command
    elif cmd == "/logout":
        cl.user_session.set("user_id", None)
        cl.user_session.set("awaiting_auth", True)
        await cl.Message(content="👋 Đã đăng xuất thành công!", author="Kumiko").send()
        await show_login_screen()
    
    # Help command
    elif cmd == "/help":
        help_msg = """
#### 📖 Hướng dẫn sử dụng

#### Lệnh đăng nhập/đăng ký:
- `/login <user_id>` - Đăng nhập với user_id có sẵn
- `/register <user_id>` - Đăng ký user_id mới
- `/logout` - Đăng xuất

#### Lệnh khác:
- `/help` - Hiển thị hướng dẫn này

#### Chat:
Sau khi đăng nhập, bạn có thể chat bình thường. Hệ thống sẽ trả lời và cung cấp tài liệu tham khảo.
"""
        await cl.Message(content=help_msg, author="Kumiko").send()
    
    else:
        await cl.Message(content=f"❌ Lệnh không hợp lệ: `{cmd}`\n\nGõ `/help` để xem danh sách lệnh.", author="Kumiko").send()


async def process_chat_message(query: str, user_id: str):
    """Process chat message and display response with progress updates"""
    # Show thinking message
    thinking_msg = cl.Message(content="🤔 Đang suy nghĩ...", author="Kumiko")
    await thinking_msg.send()
    
    try:
        # Tạo task để track progress
        start_time = asyncio.get_event_loop().time()
        
        # Send request to API (now async)
        result = await send_chat_message(query, user_id)
        
        # Remove thinking message
        await thinking_msg.remove()
        
        if not result["success"]:
            await cl.Message(content=f"{result['message']}", author="Kumiko").send()
            return
        
        answer = result["answer"]
        references = result["references"]
        
        # Store references in session for the button
        cl.user_session.set("last_references", references)
        
        # Create message with action button
        actions = [
            cl.Action(
                name="show_references",
                value="show",
                label="𝄢 Tài liệu tham khảo",
                description="Xem các nguồn tham khảo",
                payload={"references": references}
            )
        ]
        
        # Send answer with reference button
        await cl.Message(
            content=answer,
            actions=actions,
            author="Kumiko"
        ).send()
        
    except Exception as e:
        await thinking_msg.remove()
        await cl.Message(
            content=f"❌ Đã xảy ra lỗi không mong muốn: {str(e)}",
            author="Kumiko"
        ).send()


@cl.action_callback("show_references")
async def on_show_references(action: cl.Action):
    """Handle reference button click"""
    # Get references from action payload
    references = action.payload.get("references", "Không có tài liệu tham khảo")
    
    # Send references as a new message
    await cl.Message(
        content=references,
        author="Kumiko"
    ).send()
    
    # Optional: Remove the action button after clicking
    await action.remove()


@cl.on_chat_end
async def end():
    """Handle chat end"""
    print("Chat session ended")


# ============================================
# CONFIGURATION
# ============================================

@cl.set_starters
async def set_starters():
    """Set starter prompts for quick access"""
    return [
        cl.Starter(
            label="Đăng nhập",
            message="/login your_user_id",
            icon="🔑",
        ),
        cl.Starter(
            label="Đăng ký",
            message="/register your_user_id",
            icon="✍🏻",
        ),
        cl.Starter(
            label="Trợ giúp",
            message="/help",
        ),
    ]