app.py

from flask import Flask, request, jsonify
import threading
import time
import logging
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import Application, CommandHandler, MessageHandler, CallbackQueryHandler, ContextTypes, filters
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import firebase_admin
from firebase_admin import credentials, firestore
import random
import string
import os
import json
import re

# Set up logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# Initialize Flask app
app = Flask(__name__)

# Telegram Bot Config
TELEGRAM_TOKEN = "7742872311:AAEP4cVwTSyKJYLsBhrKfokNFVQboVddWpc"  # Replace with your token
BOT_USERNAME = "khCharacterAibot"
ADMIN_ID = "6728254471"

# Firebase Setup
if not os.path.exists("serviceAccountKey.json"):
    raise FileNotFoundError("Please upload 'serviceAccountKey.json' to the root directory")
cred = credentials.Certificate("serviceAccountKey.json")
if not firebase_admin._apps:
    firebase_admin.initialize_app(cred)
db = firestore.client()

# Load a smaller model and tokenizer
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # 1.1B parameters, ~2.2GB in FP16
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,  # Use FP16 for lower memory usage
    device_map="auto"  # Map to CPU
)
device = torch.device("cpu")  # Explicitly set to CPU

# Membership Plans
PLAN_LIMITS = {
    "Membership": {"ai": 1, "scenarios": 5},
    "Premium": {"ai": 2, "scenarios": 10},
    "Vip Premium Plus": {"ai": 3, "scenarios": 20}
}

# Firebase Data Functions
def get_user_info(user_id):
    user_ref = db.collection("users").document(user_id).collection("info").document("data")
    user_doc = user_ref.get()
    return user_doc.to_dict() if user_doc.exists else None

def set_user_info(user_id, data):
    db.collection("users").document(user_id).collection("info").document("data").set(data)

def get_user_ai(user_id, ai_name):
    ai_ref = db.collection("users").document(user_id).collection("ai").document(ai_name)
    ai_doc = ai_ref.get()
    return ai_doc.to_dict() if ai_doc.exists else None

def set_user_ai(user_id, ai_name, data):
    db.collection("users").document(user_id).collection("ai").document(ai_name).set(data)

def get_all_user_ai(user_id):
    ai_docs = db.collection("users").document(user_id).collection("ai").stream()
    return {doc.id: doc.to_dict() for doc in ai_docs}

def get_global_scenarios():
    scenarios = db.collection("scenarios").stream()
    return {doc.id: doc.to_dict() for doc in scenarios}

def set_global_scenario(scenario_name, data):
    db.collection("scenarios").document(scenario_name).set(data)

# Generate Smart AI Response with History
async def generate_response(context, chat_id, user_input, character, scenario, history=""):
    await context.bot.send_chat_action(chat_id=chat_id, action="typing")
    await context.bot.send_message(chat_id=chat_id, text="AI កំពុងគិត...")

    admin_data = db.collection("admin").document("training_data").get()
    training_prompt = "\n".join([f"Q: {item['question']} A: {item['answer']}" for item in admin_data.to_dict().get("data", [])]) if admin_data.exists else ""

    if re.search(r"សួស្តី|Hello|hello|hi|Hi", user_input, re.IGNORECASE):
        greetings = [
            f"សួស្តីលោកអ្នក! ខ្ញុំជា {character} នៅ {scenario}។ តើខ្ញុំអាចជួយអ្វីបាន?",
            f"សួស្តី! ខ្ញុំ {character} មកពី {scenario} រីករាយបានជួបលោកអ្នក!",
            f"សួស្តីអ្នក! នៅ {scenario} នេះ ខ្ញុំ {character} សប្បាយចិត្តណាស់!"
        ]
        return random.choice(greetings)

    system_prompt = f"""
    You are {character}, a highly intelligent AI fluent in Khmer, speaking with authority, charm, and deep understanding.
    You are in {scenario}, responding accurately and politely in natural Khmer.
    Strictly follow Khmer grammar rules (គោរពតាមក្បួនវេយ្យាករណ៍ខ្មែរ) and use polite language (ឧ. លោកអ្នក, ខ្ញុំ).
    Use this knowledge to answer correctly: {training_prompt}
    Conversation History (last 200 chars): {history[-200:]}
    """
    prompt = system_prompt + "\nអ្នកប្រើបានសួរ: " + user_input + "\n" + character + ": "

    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=1024)  # Reduced max_length for memory
    inputs = {key: value.to(device) for key, value in inputs.items()}

    outputs = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_length=1024,  # Reduced for faster CPU processing
        pad_token_id=tokenizer.eos_token_id,
        temperature=0.7,
        do_sample=True,
        top_k=50,
        top_p=0.95
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True).split(character + ": ")[-1].strip()

    if not any(ord(c) >= 0x1780 and ord(c) <= 0x17FF for c in response):  # Check for Khmer characters
        if "ជួយ" in user_input:
            response = "ខ្ញុំរីករាយនឹងជួយលោកអ្នក! សូមប្រាប់ខ្ញុំថា តើខ្ញុំអាចជួយអ្វីបាន?"
        elif "អរគុណ" in user_input:
            response = "អរគុណលោកអ្នកដែលបាននិយាយអញ្ចឹង! ខ្ញុំសប្បាយចិត្តបានជួយ។"
        else:
            response = f"ខ្ញុំមិនប្រាកដថាយល់ច្បាស់ទេ។ លោកអ្នកបានសួរ '{user_input}' មែនទេ? សូមបញ្ជាក់បន្ថែម!"
    
    return response

# Command Handlers
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = str(update.effective_user.id)
    user_data = get_user_info(user_id)
    if not user_data:
        user_data = {
            "name": update.effective_user.first_name,
            "username": update.effective_user.username or "",
            "history": "",
            "plan": "Membership",
            "banned": False,
            "stop_until": None,
            "message_count": 0,
            "last_active": time.time(),
            "ai_ratings": {},
            "selected_ai": None
        }
        set_user_info(user_id, user_data)
        await context.bot.send_message(ADMIN_ID, f"អ្នកប្រើថ្មី: {user_data['name']} (@{user_data['username']}, ID: {user_id})")
    if user_data.get("banned"):
        await update.message.reply_text("សូមទោស! លោកអ្នកត្រូវបានហាមឃាត់។")
        return
    await show_menu(update, context)

async def show_menu(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = str(update.effective_user.id)
    user_data = get_user_info(user_id)
    keyboard = [
        [InlineKeyboardButton("បង្កើត AI", callback_data="create_ai"),
         InlineKeyboardButton("កែ AI", callback_data="edit_ai")],
        [InlineKeyboardButton("ចែករំលែក AI", callback_data="share_ai"),
         InlineKeyboardButton("AI ពេញនិយម", callback_data="top_ai")],
        [InlineKeyboardButton("ប្តូរកម្មវិធី", callback_data="redeem"),
         InlineKeyboardButton("បង្ហាញ AI របស់ខ្ញុំ", callback_data="my_ai")],
        [InlineKeyboardButton("ស្ថានភាពខ្ញុំ", callback_data="my_status"),
         InlineKeyboardButton("គ្រប់គ្រង (Admin)", callback_data="admin_panel")],
        [InlineKeyboardButton("សេវាកម្មសមជាវ", callback_data="subscriptions")]
    ]
    ai_list = get_all_user_ai(user_id)
    if ai_list:
        ai_buttons = [InlineKeyboardButton(f"ជជែកជាមួយ {name}", callback_data=f"chat_with_{name}") for name in ai_list.keys()]
        keyboard.append(ai_buttons)
    reply_markup = InlineKeyboardMarkup(keyboard)
    await update.message.reply_text("សួស្តីលោកអ្នក! ខ្ញុំជា SokKha។ ជ្រើសរើសមុខងារ:", reply_markup=reply_markup)

async def my_status(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = str(update.effective_user.id)
    user_data = get_user_info(user_id)
    if not user_data:
        await update.message.reply_text("សូមចាប់ផ្តើមជាមួយ /start!")
        return
    
    active_users = sum(1 for user in db.collection("users").stream() if (info := get_user_info(user.id)) and info.get("last_active", 0) > time.time() - 3600)
    status = "សកម្ម" if user_data.get("last_active", 0) > time.time() - 3600 else "អសកម្ម"
    message_count = user_data.get("message_count", 0)
    plan = user_data.get("plan", "Membership")
    
    keyboard = [[InlineKeyboardButton(f"វាយតម្លៃ AI: {name}", callback_data=f"rate_ai_{name}") for name in get_all_user_ai(user_id).keys()]] if get_all_user_ai(user_id) else []
    reply_markup = InlineKeyboardMarkup(keyboard)
    
    await update.message.reply_text(
        f"ស្ថានភាពរបស់លោកអ្នក:\n"
        f"- ផែនការ: {plan}\n"
        f"- សកម្ម: {status}\n"
        f"- អ្នកប្រើសកម្មសរុប: {active_users}\n"
        f"- សារសរុប: {message_count}\n"
        f"ជ្រើសរើស AI ដើម្បីវាយតម្លៃ:",
        reply_markup=reply_markup
    )

# Subscription Services Display
async def show_subscriptions(update: Update, context: ContextTypes.DEFAULT_TYPE):
    subscriptions = [
        {"name": "Spud", "category": "Meal Kits", "price": "Starting at $60/week", "desc": "Fresh, local ingredient meal kits."},
        {"name": "Fubo", "category": "Sports Streaming", "price": "$79.99/month", "desc": "Live sports like NFL, NBA."},
        {"name": "CrateChef", "category": "Cooking Boxes", "price": "$49 bimonthly + shipping", "desc": "Chef-curated gourmet boxes."},
        {"name": "Wantable", "category": "Clothing Styling", "price": "$20/styling fee", "desc": "Personalized clothing boxes."}
    ]
    response = "សេវាកម្មសមជាវនៅសហរដ្ឋអាមេរិក (មិនមាននៅចក្រភពអង់គ្លេស):\n\n"
    for sub in subscriptions:
        response += f"- **{sub['name']}** ({sub['category']}): {sub['price']}\n  {sub['desc']}\n"
    await update.callback_query.edit_message_text(response)

# Button Handler
async def button(update: Update, context: ContextTypes.DEFAULT_TYPE):
    query = update.callback_query
    await query.answer()
    user_id = str(query.from_user.id)
    user_data = get_user_info(user_id)

    if not user_data:
        await query.edit_message_text("សូមចាប់ផ្តើមជាមួយ /start!")
        return
    if user_data.get("banned"):
        await query.edit_message_text("សូមទោស! លោកអ្នកត្រូវបានហាមឃាត់។")
        return

    if query.data == "create_ai":
        plan = user_data.get("plan", "Membership")
        limit = PLAN_LIMITS[plan]["ai"]
        if len(get_all_user_ai(user_id)) >= limit:
            await query.edit_message_text(f"លោកអ្នកបានឈានដល់កំណត់ {limit} AI សម្រាប់ផែនការ {plan}!")
        else:
            await query.edit_message_text("បញ្ចូលឈ្មោះ AI (ឧ. SokKha):")
            context.user_data["step"] = "create_name"
    elif query.data == "edit_ai":
        ai_list = "\n".join([name for name in get_all_user_ai(user_id).keys()]) or "គ្មាន AI"
        await query.edit_message_text(f"AI របស់លោកអ្នក:\n{ai_list}\nបញ្ចូលឈ្មោះ AI ដែលចង់កែ:")
        context.user_data["step"] = "edit_ai_name"
    elif query.data == "share_ai":
        ai_list = "\n".join([name for name in get_all_user_ai(user_id).keys()]) or "គ្មាន AI"
        await query.edit_message_text(f"AI របស់លោកអ្នក:\n{ai_list}\nបញ្ចូលឈ្មោះ AI ដែលចង់ចែករំលែក:")
        context.user_data["step"] = "share_ai"
    elif query.data == "top_ai":
        scenarios = get_global_scenarios()
        top_scenarios = sorted(scenarios.items(), key=lambda x: x[1]["usage_count"], reverse=True)[:3]
        response = "AI ពេញនិយម:\n" + "\n".join([f"- {name}: {info['usage_count']} ដង" for name, info in top_scenarios])
        await query.edit_message_text(response)
    elif query.data.startswith("chat_with_"):
        ai_name = query.data.split("chat_with_")[1]
        ai_data = get_user_ai(user_id, ai_name) or get_global_scenarios().get(ai_name)
        if not ai_data:
            await query.edit_message_text("រក AI មិនឃើញ!")
        else:
            user_data["selected_ai"] = ai_name
            set_user_info(user_id, user_data)
            response = await generate_response(context, query.message.chat_id, "សួស្តី", ai_name, ai_data["prompt"], user_data["history"])
            await query.edit_message_text(response)
            user_data["history"] += f"អ្នកប្រើ: សួស្តី\nAI: {response}\n"
            set_user_info(user_id, user_data)
    elif query.data == "redeem":
        await query.edit_message_text("បញ្ចូលកូដប័ណ្ច:")
        context.user_data["step"] = "redeem_coupon"
    elif query.data == "my_ai":
        ai_list = "\n".join([f"- {name}: {info['prompt']}" for name, info in get_all_user_ai(user_id).items()]) or "គ្មាន AI"
        await query.edit_message_text(f"AI របស់លោកអ្នក:\n{ai_list}")
    elif query.data == "my_status":
        await my_status(update, context)
    elif query.data.startswith("rate_ai_"):
        ai_name = query.data.split("rate_ai_")[1]
        keyboard = [
            [InlineKeyboardButton("⭐ 1", callback_data=f"rate_{ai_name}_1"),
             InlineKeyboardButton("⭐ 2", callback_data=f"rate_{ai_name}_2")],
            [InlineKeyboardButton("⭐ 3", callback_data=f"rate_{ai_name}_3"),
             InlineKeyboardButton("⭐ 4", callback_data=f"rate_{ai_name}_4")],
            [InlineKeyboardButton("⭐ 5", callback_data=f"rate_{ai_name}_5")]
        ]
        reply_markup = InlineKeyboardMarkup(keyboard)
        await query.edit_message_text(f"វាយតម្លៃ AI '{ai_name}':", reply_markup=reply_markup)
    elif query.data.startswith("rate_"):
        parts = query.data.split("_")
        ai_name, rating = parts[1], int(parts[2])
        user_data["ai_ratings"][ai_name] = rating
        set_user_info(user_id, user_data)
        await query.edit_message_text(f"អ្នកបានវាយតម្លៃ '{ai_name}' ជា {rating} ⭐!")
    elif query.data == "admin_panel":
        if user_id != ADMIN_ID:
            await query.edit_message_text("លោកអ្នកមិនមានសិទ្ធិ!")
            return
        keyboard = [
            [InlineKeyboardButton("រក្សាទិន្នន័យ", callback_data="admin_save"),
             InlineKeyboardButton("ផ្សព្វផ្សាយអត្ថបទ", callback_data="admin_broadcast_text")],
            [InlineKeyboardButton("ផ្សព្វផ្សាយរូបភាព", callback_data="admin_broadcast_photo"),
             InlineKeyboardButton("ផ្សព្វផ្សាយវីដេអូ", callback_data="admin_broadcast_video")],
            [InlineKeyboardButton("ផ្សព្វផ្សាយបញ្ជូនបន្ត", callback_data="admin_broadcast_forward"),
             InlineKeyboardButton("ហាមឃាត់អ្នកប្រើ", callback_data="admin_ban")],
            [InlineKeyboardButton("ដោះហាមឃាត់", callback_data="admin_unban"),
             InlineKeyboardButton("ផ្អាកអ្នកប្រើ", callback_data="admin_stop")],
            [InlineKeyboardButton("ចំនួនអ្នកប្រើ", callback_data="admin_count"),
             InlineKeyboardButton("បង្កើតកូដប័ណ្ច", callback_data="admin_coupon")],
            [InlineKeyboardButton("បណ្តុះបណ្តាល AI", callback_data="admin_train")]
        ]
        reply_markup = InlineKeyboardMarkup(keyboard)
        await query.edit_message_text("ជ្រើសមុខងារ Admin:", reply_markup=reply_markup)
    elif query.data == "subscriptions":
        await show_subscriptions(update, context)
    elif query.data.startswith("admin_"):
        await handle_admin_action(query, context)

# Handle Messages with Selected AI
async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = str(update.effective_user.id)
    user_data = get_user_info(user_id)
    text = update.message.text

    if not user_data:
        await update.message.reply_text("សូមចាប់ផ្តើមជាមួយ /start!")
        return
    if user_data.get("banned"):
        await update.message.reply_text("សូមទោស! លោកអ្នកត្រូវបានហាមឃាត់!")
        return

    user_data["message_count"] = user_data.get("message_count", 0) + 1
    user_data["last_active"] = time.time()
    set_user_info(user_id, user_data)

    if "step" in context.user_data:
        if context.user_data["step"] == "create_name":
            context.user_data["ai_name"] = text
            await update.message.reply_text("បញ្ចូលសេណារីយ៉ូ (ឧ. នៅក្នុងភូមិមួយ):")
            context.user_data["step"] = "create_scenario"
        elif context.user_data["step"] == "create_scenario":
            ai_data = {"prompt": text, "usage_count": 1, "creator": user_id}
            set_user_ai(user_id, context.user_data["ai_name"], ai_data)
            set_global_scenario(context.user_data["ai_name"], ai_data)
            await update.message.reply_text(f"បានបង្កើត AI '{context.user_data['ai_name']}'!")
            del context.user_data["step"]
        elif context.user_data["step"] == "edit_ai_name":
            ai_data = get_user_ai(user_id, text)
            if not ai_data:
                await update.message.reply_text("រក AI មិនឃើញ!")
            else:
                context.user_data["ai_name"] = text
                await update.message.reply_text("បញ្ចូលសេណារីយ៉ូថ្មី:")
                context.user_data["step"] = "edit_ai_scenario"
        elif context.user_data["step"] == "edit_ai_scenario":
            ai_data = get_user_ai(user_id, context.user_data["ai_name"])
            ai_data["prompt"] = text
            set_user_ai(user_id, context.user_data["ai_name"], ai_data)
            set_global_scenario(context.user_data["ai_name"], ai_data)
            await update.message.reply_text(f"បានកែ AI '{context.user_data['ai_name']}'!")
            del context.user_data["step"]
        elif context.user_data["step"] == "share_ai":
            if get_user_ai(user_id, text):
                share_link = f"https://t.me/{BOT_USERNAME}?start=share_{user_id}_{text}"
                await update.message.reply_text(f"លីងចែករំលែក AI '{text}': {share_link}")
            else:
                await update.message.reply_text("រក AI មិនឃើញ!")
            del context.user_data["step"]
        elif context.user_data["step"] == "redeem_coupon":
            await redeem_coupon(update, context, text)
            del context.user_data["step"]
        elif context.user_data["step"].startswith("admin_"):
            await handle_admin_input(update, context, text)
        return

    selected_ai = user_data.get("selected_ai")
    if selected_ai and (ai_data := get_user_ai(user_id, selected_ai) or get_global_scenarios().get(selected_ai)):
        response = await generate_response(context, update.message.chat_id, text, selected_ai, ai_data["prompt"], user_data["history"])
        ai_data["usage_count"] = ai_data.get("usage_count", 0) + 1
        set_user_ai(user_id, selected_ai, ai_data) if get_user_ai(user_id, selected_ai) else set_global_scenario(selected_ai, ai_data)
    else:
        response = await generate_response(context, update.message.chat_id, text, "SokKha", "នៅក្នុងភូមិមួយ", user_data["history"])
    
    await update.message.reply_text(response)
    user_data["history"] += f"អ្នកប្រើ: {text}\nAI: {response}\n"
    if len(user_data["history"]) > 1000:
        user_data["history"] = user_data["history"][-1000:]
    set_user_info(user_id, user_data)

# Admin Action Handler
async def handle_admin_action(query, context: ContextTypes.DEFAULT_TYPE):
    user_id = str(query.from_user.id)
    if user_id != ADMIN_ID:
        await query.edit_message_text("លោកអ្នកមិនមានសិទ្ធិ!")
        return

    action = query.data
    if action == "admin_save":
        users = db.collection("users").stream()
        data = {user.id: {"info": get_user_info(user.id), "ai": get_all_user_ai(user.id)} for user in users}
        with open("user_data_backup.json", "w") as f:
            json.dump(data, f, ensure_ascii=False)
        await query.edit_message_text("បានរក្សាទិន្នន័យ!")
    elif action == "admin_broadcast_text":
        await query.edit_message_text("ផ្ញើអត្ថបទ:")
        context.user_data["step"] = "admin_broadcast_text"
    elif action == "admin_broadcast_photo":
        await query.edit_message_text("ផ្ញើរូបភាព:")
        context.user_data["step"] = "admin_broadcast_photo"
    elif action == "admin_broadcast_video":
        await query.edit_message_text("ផ្ញើវីដេអូ:")
        context.user_data["step"] = "admin_broadcast_video"
    elif action == "admin_broadcast_forward":
        await query.edit_message_text("បញ្ជូនសារដែលចង់ផ្សព្វផ្សាយ:")
        context.user_data["step"] = "admin_broadcast_forward"
    elif action == "admin_ban":
        await query.edit_message_text("បញ្ចូល ID:")
        context.user_data["step"] = "admin_ban"
    elif action == "admin_unban":
        await query.edit_message_text("បញ្ចូល ID:")
        context.user_data["step"] = "admin_unban"
    elif action == "admin_stop":
        await query.edit_message_text("បញ្ចូល ID និងម៉ោង (ឧ. 12345 24):")
        context.user_data["step"] = "admin_stop"
    elif action == "admin_count":
        users = db.collection("users").stream()
        total = sum(1 for _ in users)
        await query.edit_message_text(f"ចំនួនអ្នកប្រើ: {total}")
    elif action == "admin_coupon":
        await query.edit_message_text("បញ្ចូលផែនការ និងថ្ងៃ (ឧ. Premium 30):")
        context.user_data["step"] = "admin_coupon"
    elif action == "admin_train":
        await query.edit_message_text("បញ្ចូលសំណួរ (Q):")
        context.user_data["step"] = "admin_train_q"

# Admin Input Handler
async def handle_admin_input(update: Update, context: ContextTypes.DEFAULT_TYPE, text):
    user_id = str(update.effective_user.id)
    if user_id != ADMIN_ID or "step" not in context.user_data:
        return
    step = context.user_data["step"]
    users = db.collection("users").stream()

    if step == "admin_broadcast_text":
        for user in users:
            try:
                await context.bot.send_message(user.id, text)
            except:
                pass
        await update.message.reply_text("បានផ្សព្វផ្សាយអត្ថបទ!")
    elif step == "admin_broadcast_photo" and update.message.photo:
        photo = update.message.photo[-1].file_id
        caption = update.message.caption or ""
        for user in users:
            try:
                await context.bot.send_photo(user.id, photo, caption=caption)
            except:
                pass
        await update.message.reply_text("បានផ្សព្វផ្សាយរូបភាព!")
    elif step == "admin_broadcast_video" and update.message.video:
        video = update.message.video.file_id
        caption = update.message.caption or ""
        for user in users:
            try:
                await context.bot.send_video(user.id, video, caption=caption)
            except:
                pass
        await update.message.reply_text("បានផ្សព្វផ្សាយវីដេអូ!")
    elif step == "admin_broadcast_forward":
        message_id = update.message.message_id
        for user in users:
            try:
                await context.bot.forward_message(user.id, user_id, message_id)
            except:
                pass
        await update.message.reply_text("បានបញ្ជូនបន្ត!")
    elif step == "admin_ban":
        target_id = text
        user_data = get_user_info(target_id)
        if user_data:
            user_data["banned"] = True
            set_user_info(target_id, user_data)
            await update.message.reply_text(f"បានហាមឃាត់ ID {target_id}!")
    elif step == "admin_unban":
        target_id = text
        user_data = get_user_info(target_id)
        if user_data:
            user_data["banned"] = False
            set_user_info(target_id, user_data)
            await update.message.reply_text(f"បានដោះហាមឃាត់ ID {target_id}!")
    elif step == "admin_stop":
        parts = text.split()
        if len(parts) == 2:
            target_id, hours = parts[0], int(parts[1])
            user_data = get_user_info(target_id)
            if user_data:
                user_data["stop_until"] = time.time() + hours * 3600
                set_user_info(target_id, user_data)
                await update.message.reply_text(f"បានផ្អាក ID {target_id} រយៈពេល {hours} ម៉ោង!")
    elif step == "admin_coupon":
        parts = text.split()
        if len(parts) == 2 and parts[0] in PLAN_LIMITS:
            plan, days = parts[0], int(parts[1])
            code = ''.join(random.choices(string.ascii_uppercase + string.digits, k=10))
            db.collection("coupons").document(code).set({
                "plan": plan,
                "days": days,
                "used": False,
                "created": time.time()
            })
            await update.message.reply_text(f"កូដប័ណ្ច: {code} ({plan}, {days} ថ្ងៃ)")
    elif step == "admin_train_q":
        context.user_data["train_question"] = text
        await update.message.reply_text("បញ្ចូលចម្លើយ (A):")
        context.user_data["step"] = "admin_train_a"
    elif step == "admin_train_a":
        admin_data = db.collection("admin").document("training_data").get()
        training_data = admin_data.to_dict().get("data", []) if admin_data.exists else []
        training_data.append({"question": context.user_data["train_question"], "answer": text})
        db.collection("admin").document("training_data").set({"data": training_data})
        await update.message.reply_text("បានបណ្តុះបណ្តាល AI!")
        del context.user_data["train_question"]
    del context.user_data["step"]

# Coupon Redemption
async def redeem_coupon(update: Update, context: ContextTypes.DEFAULT_TYPE, code):
    coupon = db.collection("coupons").document(code).get()
    if not coupon.exists or coupon.to_dict()["used"]:
        await update.message.reply_text("កូដមិនត្រឹមត្រូវ ឬបានប្រើរួច!")
        return
    coupon_data = coupon.to_dict()
    user_id = str(update.effective_user.id)
    user_data = get_user_info(user_id)
    user_data["plan"] = coupon_data["plan"]
    set_user_info(user_id, user_data)
    db.collection("coupons").document(code).update({"used": True, "user_id": user_id})
    await update.message.reply_text(f"បានដំឡើងផែនការ {coupon_data['plan']} សម្រាប់ {coupon_data['days']} ថ្ងៃ!")

# Application Setup
application = Application.builder().token(TELEGRAM_TOKEN).build()
application.add_handler(CommandHandler("start", start))
application.add_handler(CommandHandler("menu", show_menu))
application.add_handler(CommandHandler("mystatus", my_status))
application.add_handler(CallbackQueryHandler(button))
application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))

# Flask Routes
@app.route("/webhook", methods=["POST"])
def webhook():
    update = Update.de_json(request.get_json(force=True), application.bot)
    application.process_update(update)
    return "OK", 200

@app.route("/", methods=["GET"])
def setup_webhook():
    # Replace with your Hugging Face Space URL once known
    space_url = os.getenv("SPACE_URL", "https://seyhalite-aihavin.hf.space")
    webhook_url = f"{space_url}/webhook"
    application.bot.set_webhook(webhook_url)
    return f"Webhook set to {webhook_url}!", 200

# Run Flask in a thread
def run_flask():
    app.run(host="0.0.0.0", port=7860)

if __name__ == "__main__":
    flask_thread = threading.Thread(target=run_flask)
    flask_thread.start()
    # Start polling in the main thread
    application.run_polling()