From 8af1661b44813b88c0878d5ea4703c12c9b89a60 Mon Sep 17 00:00:00 2001
From: Azalea <22280294+hykilpikonna@users.noreply.github.com>
Date: Thu, 12 Mar 2026 00:52:09 -0400
Subject: [PATCH] =?UTF-8?q?[+]=20=E6=B5=87=E6=B0=B4=E3=80=81=E6=88=90?=
 =?UTF-8?q?=E4=B8=BA=E6=A0=91=E5=8F=B6=20buttons?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 bot.py | 41 +++++++++++++++++++++++++++++++++++++----
 db.py  | 33 +++++++++++++++++++++++++++++++--
 2 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/bot.py b/bot.py
index 4959c36..4389f4e 100644
--- a/bot.py
+++ b/bot.py
@@ -9,8 +9,8 @@ from hypy_utils import ensure_dir
 from hypy_utils.logging_utils import setup_logger
 from starlette.responses import HTMLResponse
 from starlette.staticfiles import StaticFiles
-from telegram import Update
-from telegram.ext import Application, CommandHandler, ContextTypes
+from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update
+from telegram.ext import Application, CallbackQueryHandler, CommandHandler, ContextTypes
 
 import db
 import utils
@@ -90,13 +90,18 @@ async def plant(update: Update, context: ContextTypes.DEFAULT_TYPE):
         height = db.register(channel, info.title, parent)
         await update.message.reply_text(f"""频道 {channel} 上树成功！把下面这条转发到频道里吧~""".strip())
         url_enc = urllib.parse.quote_plus(f"https://tree.aza.moe/c/{channel}")
+        leaf_text = urllib.parse.quote(f"/leaf {channel} ")
+        leaf_btn = InlineKeyboardMarkup([
+            [InlineKeyboardButton("🌿 成为树叶", url=f"https://t.me/{BOT_NAME}?text={leaf_text}")],
+            [InlineKeyboardButton("💧 浇水", callback_data=f"water:{channel}")],
+        ])
         return await update.message.reply_html(f"""
 今天是植树节，想试试和大家一起种一颗 tgcn 频道树 🌳 qwq
 
 这里是 {info.title}，是 @{parent} 的树枝 🌿 在频道树的第 {height + 1} 层哦~
 
-（如果你也有公开频道，想成为这个频道的树叶的话，就去给 @tgtreebot 发送 <code>/leaf {channel} {{你的频道名}}</code> 吧！ &gt; &lt;） <a href="https://t.me/iv?url={url_enc}&rhash=d96b84e483dc30">\u200e</a>
-""".strip())
+（如果你也有公开频道，想成为这个频道的树叶的话，就点击下面的「成为树叶」吧！ &gt; &lt;） <a href="https://t.me/iv?url={url_enc}&rhash=d96b84e483dc30">\u200e</a>
+""".strip(), reply_markup=leaf_btn)
 
     if sha not in validating:
         logger.info(f"> Channel not validated, asking for validation.")
@@ -113,9 +118,37 @@ async def plant(update: Update, context: ContextTypes.DEFAULT_TYPE):
         await update.message.reply_text("（看了一下好像频道信息还没有更新的样子... 确定加上了吗？再试试吧）")
 
 
+async def water_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    """Handle the 浇水 (water/upvote) button press."""
+    query = update.callback_query
+    data = query.data
+
+    if not data.startswith("water:"):
+        return
+
+    channel = data[len("water:"):]
+    user_id = query.from_user.id
+
+    logger.info(f"💧 Water from {user_id} {query.from_user.username or ''} for {channel}")
+
+    # Check if the channel exists
+    info = db.channel_info(channel)
+    if not info:
+        return await query.answer("这个频道不在树上哦~", show_alert=False)
+
+    # Try to add the vote
+    if db.add_vote(user_id, channel):
+        votes = db.get_votes(channel)
+        await query.answer(f"💧 浇水成功！这个树枝已经被浇了 {votes} 次水~", show_alert=False)
+    else:
+        votes = db.get_votes(channel)
+        await query.answer(f"你已经浇过水了哦~ 这个树枝已经被浇了 {votes} 次水~", show_alert=False)
+
+
 # Add handlers
 bot.add_handler(CommandHandler("start", start))
 bot.add_handler(CommandHandler("leaf", plant))
+bot.add_handler(CallbackQueryHandler(water_callback, pattern=r"^water:"))
 
 
 @app.on_event("startup")
diff --git a/db.py b/db.py
index 3cd48d4..a5e1f2b 100644
--- a/db.py
+++ b/db.py
@@ -1,4 +1,4 @@
-from peewee import Model, CharField, ForeignKeyField, IntegerField, PostgresqlDatabase
+from peewee import Model, CharField, ForeignKeyField, IntegerField, BigIntegerField, CompositeKey, PostgresqlDatabase
 
 from utils import CONFIG
 
@@ -19,8 +19,16 @@ class Channel(BaseModel):
     height = IntegerField(default=0)  # Tree height (depth)
 
 
+class Vote(BaseModel):
+    user_id = BigIntegerField()  # Telegram user ID
+    channel = ForeignKeyField(Channel, backref='votes', on_delete='CASCADE', field='username')
+
+    class Meta:
+        primary_key = CompositeKey('user_id', 'channel')
+
+
 with db:
-    db.create_tables([Channel])
+    db.create_tables([Channel, Vote])
 
 
 def channel_info(username: str) -> Channel | None:
@@ -47,6 +55,27 @@ def register(username: str, name: str, parent_username: str = None):
     return height
 
 
+def add_vote(user_id: int, channel_username: str) -> bool:
+    """Add a vote for a channel. Returns True if the vote was added, False if already voted."""
+    try:
+        Vote.create(user_id=user_id, channel=channel_username)
+        return True
+    except Exception:
+        return False
+
+
+def get_votes(channel_username: str) -> int:
+    """Get the total number of votes for a channel."""
+    return Vote.select().where(Vote.channel == channel_username).count()
+
+
+def has_voted(user_id: int, channel_username: str) -> bool:
+    """Check if a user has already voted for a channel."""
+    return Vote.select().where(
+        (Vote.user_id == user_id) & (Vote.channel == channel_username)
+    ).exists()
+
+
 if __name__ == '__main__':
     with db:
         # db.drop_tables([Channel])