import json
from json import JSONDecodeError

from database.database import get_session
from database.tvideoscript.video_script import get_today_video_script, update_video_script
from llm.local.ollama import Ollama
from log.log_manager import log
from task.manager_task import execute_task


def ai_script_task():
    with get_session() as db:
        # 1. 获取今日的热点话题列表
        video_scripts = get_today_video_script(db)
        if len(video_scripts) == 0:
            log("ai_script_task finish, task size 0")
            return
        log(f"ai script task size {len(video_scripts)}")

        ollama = Ollama()
        if not ollama.is_service_running():
            log("ai_script_task finish, ollama service not running")
            return

        for video_script in video_scripts:
            topic = video_script.title
            log(f'generate script for topic: {topic}')
            # 2. 获取话题内容
            content = video_script.content
            input_message = (
"""
## 角色
- 你是一个资深编辑。

## 目标
- 从输入的素材中选取有用的信息。

## 任务描述
- 从contents中选取最能吸引人的段落或句子，使读者产生兴趣和共鸣。请确保这些内容具有情感张力、戏剧性、趣味性或引发思考的价值。选取两部分内容，分别使用键“content_one”和“content_two”。每部分不少于100个汉字。两部分内容字数之和不多于600个汉字。

## 要求
- 严格遵守字数要求。
- 直接输出内容。
- 内容为JSON格式。

## 素材如下
""" + content
            )
            # log(input_message)
            # 3. 调用ollama生成话题脚本
            llm_result = ollama.generate_text(input_message)
            log(llm_result)
            try:
                llm_result = json.loads(llm_result)
            except JSONDecodeError as e:
                log(f"ai_script_task error: {e}. skip topic: {topic}")
                continue
            video_script.script = (
f"""{video_script.title}
{video_script.description}
一位网友说：
{llm_result["content_one"]}
另一位网友说:
{llm_result["content_two"]}
关于这个问题大家有什么看法呢？
欢迎评论区留言
"""
            )
            # 4. 保存话题脚本
            update_video_script(db, video_script)

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