import time

from database.database import get_session
from database.tcontentdispatch.curd import get_content_by_title_and_category, update
from llm.local.ollama import Ollama
from log.log_manager import log
from task.manager_task import execute_task


def ai_edit_with_ollama(input_message: str) -> str:
    log(f"ai_edit_with_ollama start execute at {time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())}")
    ollama = Ollama()
    response = ollama.generate_text(input_message)
    log(f"ai_edit_with_ollama end execute at {time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())}")
    return response

def revisal_task():
    with get_session() as db:
        # 1. 构建楼市参考消息文章标题，格式：楼市参考消息yyyy-MM-dd
        title = ("楼市参考消息" +
                 time.strftime("%Y", time.localtime()) + '年' +
                 time.strftime("%m", time.localtime()) + '月' +
                 time.strftime("%d", time.localtime()) + '日')
        # 2. 从内容分发数据表获取当前标题和分类的文章是否存在
        content_dispatch = get_content_by_title_and_category(db, title, "房地产")
        ai_content = ""
        if content_dispatch and content_dispatch.content:
            print(content_dispatch.content)
            # 3. 执行AI编辑
            input_message = (('按照规则编辑提供的内容。规则如下：\n'
                             '1 以每行内容作为一个处理单元，去掉与房地产、楼市无关的内容。如果整行内容都无关则全部去掉，否则全部保留。\n'
                             '2 去除重复内容\n'
                             '3 不要故意删除内容\n'
                             '5 不要出现空行\n'
                             '6 重新编号\n'
                             '7 不要出现类似"以下是根据您提供的规则编辑后的内容"等提示信息，直接输出编辑后的内容\n'
                             '内容如下：\n')
                             + content_dispatch.content)
            ai_content = ai_edit_with_ollama(input_message)
            print("-----------------------------------------------------------")
            print(ai_content)
        # 4. 把content写入数据库
        if ai_content:
            content_dispatch.ai_content = ai_content
            content_dispatch.is_sent = False
            update(db)


if __name__ == "__main__":
    execute_task(revisal_task)