from dotenv import load_dotenv
load_dotenv(verbose=True)

import os
import time
import signal
import shutil
import json
import traceback
from datetime import datetime, timedelta

print('✨ 正在初始化 INF-SJTU ...')

from modules.core.spider_manage import spmanage
from modules.core.news_manage import newsmanage

# CONFIG
from config import *

# MAIN
class Worker:
    def __init__(self) -> None:
        
        # 根目录
        self.root_path = os.path.dirname(os.path.abspath(__file__))
        
        # 检测缓存目录
        os.makedirs(f'{self.root_path}/cache', exist_ok=True)
        self.get_cache_path = lambda path: os.path.join(self.root_path, f'cache/{path}')
        self.get_data_path = lambda path: os.path.join(self.root_path, f'data/{path}')
        
        # 当前日期
        self.current_date = datetime.now().strftime("%Y-%m-%d")
        
        # Modules
        # 爬虫
        self.spmanage = spmanage(self.root_path, self.get_cache_path, self.current_date)
        # 新闻管理
        self.newsmanage = newsmanage(self.root_path, self.get_cache_path, self.current_date)
        
        pass
    
    def run(self):
        print('🧐 等待执行中！')
        count = 0
        while True:
            try:
                count += 1
                print(f'✨ 第 {count} 次任务执行中...\n')
                
                self.check_news_updated()
                self.clear_old_news()
            
                print('\n🧐 任务执行完毕！正在等待下一轮任务...')
                time.sleep(EXECUTION_CYCLE)  # 等待 1 小时
                
            # 截获异常
            except:
                print('\n🚨 INF-SJTU 运行异常！\n')
                traceback.print_exc()
                time.sleep(1)
                continue
        
    """ MAIN | 定时触发 """
    def check_news_updated(self):
        """ 检测是否有新的新闻 """
        print('🔖 正在检测是否有新的新闻...')
        
        # 检测缓存目录中是否有对应日期的新闻
        # 解放日报
        jfdaily_news_path = self.get_cache_path(f'news/{self.current_date}/jfdaily/news.json')
        if not os.path.exists(jfdaily_news_path):
            # 触发爬取
            raw_news = self.spmanage.jfdaily()
            if raw_news:
                # 预处理新闻
                self.newsmanage.jfdaily_news_convert(raw_news)
                # 新闻分类
                self.newsmanage.news_classify('jfdaily')
                # 索引新闻
                # self.newsmanage.news_index('jfdaily')
                # 生成问题
                self.newsmanage.generate_qa('jfdaily')
                # 阅读分级
                self.newsmanage.generate_reading_level('jfdaily')
        else:
            # 检测补全缺少的信息
            print('🔖 检测补全缺少的信息...')
            with open(jfdaily_news_path, 'r', encoding='utf-8') as f:
                raw_news = json.loads(f.read())
                for news in raw_news:
                    # 检查 questions 是否存在
                    if 'questions' not in raw_news[news]:
                        print(f'🚧 侦测到「解放日报」的 {news} 新闻缺少问题！正在补全...')
                        # 生成问题
                        self.newsmanage.generate_qa('jfdaily', news)
                    
                    # 检查 reading_level 是否存在
                    if 'reading_level' not in raw_news[news]:
                        print(f'🚧 侦测到「解放日报」的 {news} 新闻缺少阅读分级！正在补全...')
                        # 阅读分级
                        self.newsmanage.generate_reading_level('jfdaily')

        # Nature Brief
        naturebrief_news_path = self.get_cache_path(f'news/{self.current_date}/naturebrief/news.json')
        if not os.path.exists(naturebrief_news_path):
            # 触发爬取
            raw_news = self.spmanage.naturebrief()
            if raw_news:
                # 预处理新闻
                self.newsmanage.naturebrief_news_convert(raw_news)
                # 新闻分类
                self.newsmanage.news_classify('naturebrief')
                # 索引新闻
                # self.newsmanage.news_index('naturebrief')
                # 生成问题
                self.newsmanage.generate_qa('naturebrief')
                # 阅读分级
                self.newsmanage.generate_reading_level('naturebrief')
        else:
            # 检测补全缺少的信息
            with open(naturebrief_news_path, 'r', encoding='utf-8') as f:
                raw_news = json.loads(f.read())
                for news in raw_news:
                    # 检查 questions 是否存在
                    if 'questions' not in raw_news[news]:
                        print(f'🚧 侦测到「Nature Brief」的 {news} 新闻缺少问题！正在补全...')
                        # 生成问题
                        self.newsmanage.generate_qa('naturebrief')
                    
                    # 检查 reading_level 是否存在
                    if 'reading_level' not in raw_news[news]:
                        print(f'🚧 侦测到「Nature Brief」的 {news} 新闻缺少阅读分级！正在补全...')
                        # 阅读分级
                        self.newsmanage.generate_reading_level('naturebrief')
            
        print('🔖 检测完毕！')
    
    def clear_old_news(self):
        """ 清理缓存目录下旧的新闻 | 三天前 """
        print('🔖 正在清理缓存目录下旧的新闻...')
        
        # 校验新闻目录
        if not os.path.exists(self.get_cache_path('news')):
            os.makedirs(self.get_cache_path(f'news/{self.current_date}'))
            return
        
        # 获取缓存目录下的第一级目录名列表
        cache_dir_list = os.listdir(self.get_cache_path('news'))
        
        # 获取近三天的日期列表
        date_list = []
        for i in range(3):
            date_list.append((datetime.now() - timedelta(days=i)).strftime("%Y-%m-%d"))
        
        # 清理
        count = 0
        for cache_dir in cache_dir_list:
            if cache_dir not in date_list:
                del_dir_path = self.get_cache_path(f'news/{cache_dir}')
                
                # 解放日报向量数据库映射
                jfdaily_db_mapping_path = self.get_cache_path(f'news/{cache_dir}/jfdaily/db_mapping.json')
                # 删除对应的向量数据库
                # 解析映射 json
                with open(jfdaily_db_mapping_path, 'r', encoding='utf-8') as f:
                    db_mappings = json.loads(f.read())
                    # 删除对应的向量数据库
                    for db_mapping in db_mappings:
                        shutil.rmtree(self.get_data_path(f'vectordb/{db_mapping["dbName"]}'))
                
                # # 递归删除缓存
                shutil.rmtree(del_dir_path)
                count += 1
        
        if count == 0:
            print('🔖 无旧新闻可清理！')
        else:
            print(f'🔖 清理完毕！共清理 {count} 天的旧新闻')        
                
# 调试模式 | 直接执行
if __name__ == "__main__":
    # 终止信号处理
    def sigterm_handler(_signo, _stack_frame):
        print('\n🚨 INF-SJTU 退出运行！')
        os._exit(0)

    signal.signal(signal.SIGTERM, sigterm_handler)
    signal.signal(signal.SIGINT, sigterm_handler)
    
    # 清空终端日志
    os.system('cls||clear')
    print('🐵 INF-SJTU RUNNING...\n')
    
    # 启动 Worker
    worker = Worker()
    worker.run()
File "f:\tpyrced_inf-delivery-sjtu-main\inf-delivery-sjtu-main\main.py", line 1, in <module>
    from dotenv import load_dotenv
ModuleNotFoundError: No module named 'dotenv'
怎么解决？