src/api.py

from fastapi import FastAPI, Request
from fastapi.responses import RedirectResponse
import httpx
import asyncio
import getpass
import gradio as gr
import hashlib
import bcrypt
import re
from src.bot import Bot
from src.webhook import WebhookUpdate
from src.logger import logger
from config.settings import Settings
import time
import os
import socket
from typing import List, Dict
from gradio import mount_gradio_app
from datetime import datetime
from pydantic import BaseModel

app = FastAPI()
http_client = httpx.AsyncClient()
app.state.bot_lock = asyncio.Lock()

# Route chuyển hướng từ / đến /gradio
@app.get("/")
async def redirect_to_gradio():
    logger.info("[redirect] Redirecting from / to /gradio")
    return RedirectResponse(url="/gradio")

# Endpoint /chat cho tích hợp API (tùy chọn)
class ChatRequest(BaseModel):
    prompt: str
    user_id: str = "4096249"
    is_gradio: bool = False

@app.post("/chat")
async def chat_endpoint(request: ChatRequest):
    try:
        chat_id = int(hashlib.sha256((str(request.user_id) + str(int(time.time()))).encode()).hexdigest(), 16) % 10**6
        response = await app.state.bot.handle_message(chat_id, request.user_id, request.prompt, request.is_gradio)
        response_str = "<br>".join(str(item).strip() for item in response if item) if isinstance(response, list) else str(response).strip()
        return {"response": response_str}
    except Exception as e:
        logger.error(f"[chat_endpoint] Error: {str(e)}")
        return {"error": str(e)}

def is_port_in_use(port):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        return s.connect_ex(('localhost', port)) == 0

async def chatbot_handler(message: str, history: List[Dict], user_id: int = None) -> List[Dict]:
    logger.info(f"[chatbot_handler] Start processing message: {message[:50]}...")
    try:
        if not message or not isinstance(message, str):
            logger.error("[chatbot_handler] Invalid message: empty or not a string")
            return history + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": "Lỗi: Tin nhắn không hợp lệ."}
            ]

        async with app.state.bot_lock:
            if not hasattr(app.state, "bot") or app.state.bot is None:
                logger.warning("[chatbot_handler] Bot not initialized. Initializing...")
                app.state.bot = Bot()
                await app.state.bot.initialize()
                await app.state.bot.cache.clear_all()

            if not user_id:
                return history + [
                    {"role": "user", "content": message},
                    {"role": "assistant", "content": "Vui lòng đăng nhập bằng email/tên người dùng và mật khẩu."}
                ]

            chat_id = int(hashlib.sha256((str(user_id) + str(int(time.time()))).encode()).hexdigest(), 16) % 10**6
            response = await app.state.bot.handle_message(chat_id, user_id, message, is_gradio=True)
            new_history = history.copy()
            new_history.append({"role": "user", "content": message.strip()})
            if isinstance(response, list):
                for part in response:
                    new_history.append({"role": "assistant", "content": str(part).strip()})
            else:
                new_history.append({"role": "assistant", "content": str(response).strip()})
            logger.info(f"[chatbot_handler] Added response to history: {new_history}")
            return new_history
    except asyncio.TimeoutError:
        logger.error("[chatbot_handler] Timeout processing message", exc_info=True)
        return history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": "Lỗi: Timeout khi xử lý tin nhắn."}
        ]
    except Exception as e:
        logger.error(f"[chatbot_handler] Error: {str(e)}", exc_info=True)
        return history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": f"Lỗi: {str(e)}"}
        ]

def create_gradio_interface():
    with gr.Blocks(
        theme=gr.themes.Soft(),
        title="CotienBot",
        css="""
            .gradio-container { max-width: 800px; margin: auto; padding: 20px; }
            #chatbot-container { max-height: 500px; min-height: 300px; overflow-y: auto; border: 1px solid #ccc; border-radius: 10px; padding: 15px; background: #f9f9f9; margin-bottom: 10px; -webkit-overflow-scrolling: touch; }
            #chatbot-container .gr-chatbot { font-family: Arial, sans-serif; font-size: 14px; }
            #chatbot-container .gr-chatbot .message.user { background: #e6f3ff; padding: 10px; margin: 5px 0; border-radius: 8px; }
            #chatbot-container .gr-chatbot .message.assistant { background: #f0f0f0; padding: 10px; margin: 5px 0; border-radius: 8px; }
            input, button { font-family: Arial, sans-serif; font-size: 14px; }
            .gr-button { margin: 5px; }
            .gr-textbox { width: 100%; box-sizing: border-box; border: 1px solid #ccc; border-radius: 5px; }
            .gr-markdown { margin-bottom: 10px; font-size: 16px; }
            #input-container { display: flex; align-items: center; margin-top: 10px; }
            #input-container .gr-textbox { flex-grow: 1; margin-right: 10px; }
            #input-container .gr-button { flex-shrink: 0; }
        """,
        head="""
            <script>
                document.addEventListener('DOMContentLoaded', function() {
                    console.log('Initializing auto-scroll for chatbot');
                    const chatbotContainer = document.getElementById('chatbot-container');
                    if (chatbotContainer) {
                        console.log('Chatbot container found');
                        const observer = new MutationObserver(function(mutations) {
                            console.log('Chatbot content changed, scrolling to bottom');
                            chatbotContainer.scrollTo({ top: chatbotContainer.scrollHeight, behavior: 'smooth' });
                        });
                        observer.observe(chatbotContainer, { childList: true, subtree: true, characterData: true });
                        chatbotContainer.addEventListener('gradio:render', function() {
                            console.log('Gradio render event triggered, scrolling to bottom');
                            chatbotContainer.scrollTo({ top: chatbotContainer.scrollHeight, behavior: 'smooth' });
                        });
                    } else {
                        console.error('Chatbot container not found');
                    }
                });
            </script>
        """
    ) as interface:
        with gr.Column():
            gr.Markdown("# CotienBot - Xổ số AI")
            gr.Markdown("Đăng ký hoặc đăng nhập để sử dụng bot. Sử dụng lệnh `/register <email/tên_người_dùng> <mật_khẩu>` để đăng ký.")
            gr.Markdown("Đối với lệnh `/load_lottery`, nhập kết quả xổ số theo định dạng nhiều dòng, ví dụ:\n```\n/load_lottery Xổ số Đắk Lắk 17-06-2025\nĐặc biệt 123456\nGiải nhất 12345\n...\n```")
            with gr.Row():
                identifier_input = gr.Textbox(label="Email hoặc Tên người dùng", placeholder="Nhập email/tên người dùng", lines=1, max_lines=1)
                password_input = gr.Textbox(label="Mật khẩu người dùng", placeholder="Nhập mật khẩu", type="password", lines=1, max_lines=1)
            with gr.Row():
                register_btn = gr.Button("Đăng ký")
                login_btn = gr.Button("Đăng nhập")
            chatbot = gr.Chatbot(label="Lịch sử trò chuyện", type="messages", height=500, elem_id="chatbot-container")
            with gr.Row(elem_id="input-container"):
                msg = gr.Textbox(
                    placeholder="Nhập lệnh (VD: /load_lottery với kết quả xổ số nhiều dòng hoặc /auth <mật_khẩu_bot>)",
                    label="Tin nhắn",
                    interactive=True,
                    lines=10,
                    max_lines=20
                )
                submit_btn = gr.Button("Gửi")
            clear_btn = gr.Button("Xóa")
            gr.Examples(
                examples=[
                    "/register user@example.com mypassword",
                    "/auth 1234",
                    "/logout",
                    "/help",
                    "/load_lottery Xổ số Bình Định Thứ năm 03-07-2025\nĐặc biệt 162010\nGiải nhất 63575\nGiải nhì 45370\nGiải ba 73648 76616\nGiải tư 13393 64399 56592 91472 82442 95757 03648\nGiải năm 5612\nGiải sáu 2310 1286 1335\nGiải bảy 417\nGiải tám 89",
                    "/lottery_position xskh 18-05-2025 3",
                    "/check_dates"
                ],
                inputs=msg
            )

        user_id_state = gr.State(value=None)

        async def handle_register(identifier, user_password, history):
            logger.info(f"[handle_register] identifier={identifier}")
            if not identifier.strip() or not user_password.strip():
                return history + [
                    {"role": "assistant", "content": "Vui lòng nhập email/tên người dùng và mật khẩu."}
                ], identifier, user_password
            try:
                command = f"/register {identifier} {user_password}"
                temp_user_id = int(hashlib.sha256(f"{identifier}:{user_password}".encode()).hexdigest(), 16) % 10000000
                result = await app.state.bot.handle_message(chat_id=0, user_id=temp_user_id, text=command, is_gradio=True)
                new_history = history.copy()
                new_history.append({"role": "user", "content": command})
                if isinstance(result, list):
                    for part in result:
                        new_history.append({"role": "assistant", "content": str(part).strip()})
                else:
                    new_history.append({"role": "assistant", "content": str(result).strip()})
                if "Đăng ký thành công" in str(result):
                    combined = f"{identifier}:{user_password}"
                    user_id = int(hashlib.sha256(combined.encode()).hexdigest(), 16) % 10000000
                logger.debug(f"[handle_register] After register, new_history: {new_history}")
                return new_history, identifier, user_password
            except Exception as e:
                logger.error(f"[handle_register] Lỗi: {str(e)}")
                return history + [
                    {"role": "assistant", "content": f"Lỗi: {str(e)}"}
                ], identifier, user_password

        async def handle_login(identifier, user_password, history):
            logger.info(f"[handle_login] identifier={identifier}")
            if not identifier.strip() or not user_password.strip():
                return None, history + [
                    {"role": "assistant", "content": "Vui lòng nhập email/tên người dùng và mật khẩu."}
                ], identifier, user_password
            try:
                combined = f"{identifier}:{user_password}"
                user_id = int(hashlib.sha256(combined.encode()).hexdigest(), 16) % 10000000
                doc_id = f"gradio_{user_id}"
                doc = await app.state.bot.db.get(data_type="users", doc_id=doc_id)
                if not doc or doc.get("type") != "gradio":
                    return None, history + [
                        {"role": "assistant", "content": "Sai email/tên người dùng hoặc mật khẩu."}
                    ], identifier, user_password
                if not bcrypt.checkpw(user_password.encode(), doc.get("password_hash").encode()):
                    return None, history + [
                        {"role": "assistant", "content": "Sai email/tên người dùng hoặc mật khẩu."}
                    ], identifier, user_password
                await app.state.bot.db.set(
                    {
                        "authenticated": True,
                        "last_auth": datetime.now().isoformat()
                    },
                    data_type="users",
                    doc_id=doc_id,
                    merge=True
                )
                new_history = history.copy()
                new_history.append({"role": "assistant", "content": f"Đăng nhập thành công! user_id={user_id}"})
                logger.debug(f"[handle_login] After login, new_history: {new_history}")
                return user_id, new_history, identifier, user_password
            except Exception as e:
                logger.error(f"[handle_login] Lỗi: {str(e)}")
                return None, history + [
                    {"role": "assistant", "content": f"Lỗi: {str(e)}"}
                ], identifier, user_password

        async def handle_submit(identifier, user_password, message, history, user_id):
            logger.info(f"[handle_submit] identifier={identifier}, message={message[:50]}..., user_id={user_id}")
            logger.debug(f"[handle_submit] Original message: {message}")
            if not message.strip():
                return "", history + [
                    {"role": "assistant", "content": "Vui lòng nhập lệnh."}
                ], user_id, identifier, user_password
            try:
                # Chỉ chuẩn hóa cho /load_lottery và /add_lottery
                normalized_message = message.strip()
                if message.lower().startswith(("/load_lottery", "/add_lottery")):
                    normalized_message = re.sub(r'\s+', ' ', message.strip())
                    for prize in ["Đặc biệt", "Giải nhất", "Giải nhì", "Giải ba", "Giải tư", "Giải năm", "Giải sáu", "Giải bảy", "Giải tám"]:
                        normalized_message = re.sub(rf'({prize}\s)', f'\n\\1', normalized_message, flags=re.IGNORECASE)
                logger.debug(f"[handle_submit] Normalized message: {normalized_message}")

                command = normalized_message.split(" ")[0].lower()
                public_commands = ["/start", "/help", "/auth", "/register", "/logout"]
                
                # Luôn thêm tin nhắn gốc của người dùng vào lịch sử
                new_history = history.copy()
                new_history.append({"role": "user", "content": message.strip()})

                if command == "/register":
                    parts = normalized_message.split()
                    if len(parts) != 3:
                        new_history.append({"role": "assistant", "content": "Cú pháp: /register <email/tên_người_dùng> <mật_khẩu>"})
                        logger.debug(f"[handle_submit] Invalid /register syntax, new_history: {new_history}")
                        return "", new_history, user_id, identifier, user_password
                    identifier = parts[1]
                    user_password = parts[2]
                    temp_user_id = int(hashlib.sha256(f"{identifier}:{user_password}".encode()).hexdigest(), 16) % 10000000
                    result = await app.state.bot.handle_message(chat_id=0, user_id=temp_user_id, text=normalized_message, is_gradio=True)
                    if isinstance(result, list):
                        for part in result:
                            new_history.append({"role": "assistant", "content": str(part).strip()})
                    else:
                        new_history.append({"role": "assistant", "content": str(result).strip()})
                    if "Đăng ký thành công" in str(result):
                        user_id = temp_user_id
                    logger.debug(f"[handle_submit] After /register, new_history: {new_history}")
                    return "", new_history, user_id, identifier, user_password

                if not user_id:
                    new_history.append({"role": "assistant", "content": "Vui lòng đăng nhập trước khi gửi lệnh."})
                    logger.debug(f"[handle_submit] No user_id, new_history: {new_history}")
                    return "", new_history, user_id, identifier, user_password

                # Kiểm tra xác thực cho lệnh nâng cao
                if command in Bot.ADVANCED_COMMANDS:
                    is_authenticated, message = await app.state.bot.check_user_auth(user_id, is_gradio=True)
                    if not is_authenticated:
                        new_history.append({"role": "assistant", "content": message})
                        logger.debug(f"[handle_submit] Auth failed, new_history: {new_history}")
                        return "", new_history, user_id, identifier, user_password

                result = await app.state.bot.handle_message(chat_id=0, user_id=user_id, text=normalized_message, is_gradio=True)
                if isinstance(result, list):
                    for part in result:
                        new_history.append({"role": "assistant", "content": str(part).strip()})
                else:
                    new_history.append({"role": "assistant", "content": str(result).strip()})
                logger.debug(f"[handle_submit] After handling message, new_history: {new_history}")
                return "", new_history, user_id, identifier, user_password
            except Exception as e:
                logger.error(f"[handle_submit] Lỗi: {str(e)}")
                new_history.append({"role": "assistant", "content": "Lỗi hệ thống, vui lòng thử lại sau."})
                logger.debug(f"[handle_submit] After error, new_history: {new_history}")
                return "", new_history, user_id, identifier, user_password

        async def handle_clear():
            return "", [], None, "", ""

        register_btn.click(
            fn=handle_register,
            inputs=[identifier_input, password_input, chatbot],
            outputs=[chatbot, identifier_input, password_input],
            queue=True
        )
        login_btn.click(
            fn=handle_login,
            inputs=[identifier_input, password_input, chatbot],
            outputs=[user_id_state, chatbot, identifier_input, password_input],
            queue=True
        )
        submit_btn.click(
            fn=handle_submit,
            inputs=[identifier_input, password_input, msg, chatbot, user_id_state],
            outputs=[msg, chatbot, user_id_state, identifier_input, password_input],
            queue=True
        )
        clear_btn.click(
            fn=handle_clear,
            outputs=[msg, chatbot, user_id_state, identifier_input, password_input]
        )
    return interface

@app.get("/health")
async def health():
    logger.info("[health] Health check")
    return {"status": "OK"}

@app.get("/check_cache_dir")
async def check_cache_dir():
    cache_dir = Settings.CACHE_DIR
    logger.info(f"[check_cache_dir] Checking cache: {cache_dir}")
    try:
        if not os.path.exists(cache_dir):
            os.makedirs(cache_dir, mode=0o777, exist_ok=True)
        return {
            "cache_dir": cache_dir,
            "writable": os.access(cache_dir, os.W_OK),
            "permissions": oct(os.stat(cache_dir).st_mode & 0o777)
        }
    except Exception as e:
        logger.error(f"[check_cache_dir] Error: {str(e)}")
        return {"error": str(e)}

@app.post("/v1/webhook")
async def webhook(request: Request):
    logger.info("[webhook] Received Telegram webhook request")
    try:
        async with asyncio.timeout(60):
            async with app.state.bot_lock:
                data = await request.json()
                update = WebhookUpdate(**data)
                if not hasattr(app.state, "bot") or not app.state.bot:
                    logger.error("[webhook] Bot not initialized")
                    return {"status": "error", "message": "Bot not initialized"}
                logger.info("[webhook] Calling handle_update")
                await app.state.bot.handle_update(update)
                return {"status": "ok"}
    except Exception as e:
        logger.error(f"[webhook] Error: {str(e)}", exc_info=True)
        return {"status": "error", "message": str(e)}

@app.on_event("startup")
async def startup_event():
    logger.info("[startup] Starting up...")
    app.state.bot = Bot()
    app.state.gradio_iface = None
    try:
        async with asyncio.timeout(120):
            cache_dir = Settings.CACHE_DIR
            logger.info(f"Khởi động với user: {getpass.getuser()}, cache_dir={cache_dir}")
            if not os.access(cache_dir, os.W_OK):
                raise RuntimeError(f"Không có quyền ghi vào {cache_dir}")
            await app.state.bot.initialize()  # Khởi tạo bot và đặt lại trạng thái xác thực
            await app.state.bot.cache.clear_all()
            webhook_url = f"{Settings.WEBHOOK_URL.rstrip('/')}/v1/webhook"
            logger.info(f"[startup] Setting webhook: {webhook_url}")
            response = await http_client.post(
                f"https://api.telegram.org/bot{Settings.TELEGRAM_TOKEN}/getWebhookInfo",
                timeout=10.0
            )
            logger.info(f"[startup] Webhook info: {response.json()}")
            if response.json().get("result", {}).get("url") != webhook_url:
                response = await http_client.post(
                    f"https://api.telegram.org/bot{Settings.TELEGRAM_TOKEN}/setWebhook",
                    json={"url": webhook_url},
                    timeout=10.0
                )
                logger.info(f"[startup] Set webhook: {response.json()}")
            app.state.gradio_iface = create_gradio_interface()
            logger.info("[startup] Gradio interface created")
    except Exception as e:
        logger.error(f"[startup] Error: {str(e)}", exc_info=True)
        raise

@app.on_event("shutdown")
async def shutdown_event():
    try:
        if app.state.bot.db and hasattr(app.state.bot.db, 'close'):
            await app.state.bot.db.close()
            logger.info("Đã đóng FirestoreDB")
        if app.state.bot.search and hasattr(app.state.bot.search, 'close'):
            await app.state.bot.search.close()
            logger.info("Đã đóng Search")
        await http_client.aclose()
        logger.info("Đã đóng http_client")
        if app.state.gradio_iface:
            app.state.gradio_iface.close()
            logger.info("Đã đóng Gradio")
        logger.info("Bot đã tắt hoàn toàn")
    except Exception as e:
        logger.error(f"[shutdown] Lỗi khi tắt bot: {str(e)}")

app = mount_gradio_app(app, create_gradio_interface(), path="/gradio")