From f504832ac53e2469dd1d90576009d844b252c71c Mon Sep 17 00:00:00 2001
From: xtaodada <xtao@xtaolink.cn>
Date: Thu, 5 Sep 2024 18:47:18 +0800
Subject: [PATCH] feat: weather graph

---
 defs/weather_graph.py    | 70 ++++++++++++++++++++++++++++++++++++++++
 modules/weather_graph.py | 40 +++++++++++++++++++++++
 2 files changed, 110 insertions(+)
 create mode 100644 defs/weather_graph.py
 create mode 100644 modules/weather_graph.py

diff --git a/defs/weather_graph.py b/defs/weather_graph.py
new file mode 100644
index 0000000..3d0cd3f
--- /dev/null
+++ b/defs/weather_graph.py
@@ -0,0 +1,70 @@
+import asyncio
+import os
+import shutil
+from pathlib import Path
+
+import aiofiles
+import httpx
+from bs4 import BeautifulSoup
+
+from defs.bilibili_download import execute
+
+FFMPEG_PATH = "ffmpeg"
+URL = "http://www.nmc.cn/publish/satellite/fy4b-visible.htm"
+TEMP_PATH = Path("data") / "weather_tmp"
+TEMP_PATH.mkdir(exist_ok=True, parents=True)
+OUTPUT_PATH = TEMP_PATH / "output.mp4"
+
+
+class WeatherGraphError(Exception):
+    def __init__(self, error: str):
+        self.error = error
+
+
+async def download_img(client, index: int, url: str):
+    try:
+        res = await client.get(url)
+        if res.status_code != 200:
+            print(f"下载天气图片失败，错误代码 {res.status_code}")
+        async with aiofiles.open(
+            TEMP_PATH / f"img{str(index).zfill(3)}.jpg", "wb"
+        ) as f:
+            await f.write(res.content)
+    except Exception as e:
+        print(f"下载天气图片失败 {url} {e}")
+
+
+async def get_images():
+    async with httpx.AsyncClient() as client:
+        res = await client.get(URL)
+        if res.status_code != 200:
+            raise WeatherGraphError(f"获取图片数据失败，错误代码 {res.status_code}")
+        html = res.content
+        soup = BeautifulSoup(html, "lxml")
+        elements = soup.find_all(class_="time")
+        tasks = []
+        for index, element in enumerate(reversed(elements)):
+            url = element["data-img"]
+            url = url.replace("medium/", "")
+            task = asyncio.create_task(download_img(client, index, url))
+            tasks.append(task)
+        await asyncio.gather(*tasks)
+
+
+async def get_video():
+    frame_rate = 6
+    ffmpeg_cmd = f'{FFMPEG_PATH} -r {frame_rate} -i "{TEMP_PATH}/img%03d.jpg" -vf scale=1920:1080 -c:v libx264 -pix_fmt yuv420p -y "{TEMP_PATH}/output.mp4"'
+    _, code = await execute(ffmpeg_cmd)
+    if code != 0:
+        raise WeatherGraphError("生成视频失败")
+
+
+async def gen() -> Path:
+    if not os.path.exists(TEMP_PATH):
+        os.mkdir(TEMP_PATH)
+    else:
+        shutil.rmtree(TEMP_PATH)
+        os.mkdir(TEMP_PATH)
+    await get_images()
+    await get_video()
+    return OUTPUT_PATH
diff --git a/modules/weather_graph.py b/modules/weather_graph.py
new file mode 100644
index 0000000..5576cb4
--- /dev/null
+++ b/modules/weather_graph.py
@@ -0,0 +1,40 @@
+from pyrogram import filters, Client
+from pyrogram.types import Message
+
+from init import bot, logger
+
+from cashews import cache
+from defs.weather_graph import OUTPUT_PATH, gen, WeatherGraphError
+from scheduler import scheduler
+
+CACHE_KEY = "weather_graph:file_id"
+
+
+@bot.on_message(
+    filters.incoming
+    & filters.command(["weather_graph", f"weather_graph@{bot.me.username}"])
+)
+async def weather_graph_command(_: "Client", message: "Message"):
+    if file_id := await cache.get(CACHE_KEY):
+        await message.reply_video(file_id, quote=True)
+        return
+    if OUTPUT_PATH.exists():
+        reply = await message.reply_video(OUTPUT_PATH, quote=True)
+        await cache.set(CACHE_KEY, reply.video.file_id, expire=3600)
+    else:
+        reply = await message.reply("正在生成中...")
+        try:
+            await gen()
+            r1 = await message.reply_video(OUTPUT_PATH, quote=True)
+            await cache.set(CACHE_KEY, r1.video.file_id, expire=3600)
+            await reply.delete()
+        except WeatherGraphError as e:
+            await reply.edit(f"生成失败：{e.error}")
+
+
+@scheduler.scheduled_job("interval", hours=1, id="weather_graph_refresh")
+async def weather_graph_refresh():
+    try:
+        await gen()
+    except WeatherGraphError as e:
+        logger.exception("Weather Graph Refresh failed", exc_info=e)