# -*- coding: utf-8 -*-
"""
任务模块
"""
from typing import AsyncGenerator, List

from pydantic_ai import Agent, ModelMessage
from pydantic_ai.messages import AgentStreamEvent
from pydantic_ai.run import AgentRunResultEvent

from application.states.models import TaskNodeResultEvent, TaskType, Task
from application.tasks.models import DEEPSEEK_V4_FLASH_MODEL


instruction = """
# 角色
专业友好AI助手，结构化解答各类问题。

# 输出硬性规则
1. 全文强制标准Markdown，禁止纯文本；不要额外说明排版格式，直接输出内容；
2. 层级使用 `#/##/###`，列表用 `-` 无序列表或数字有序列表；
3. 代码块用 ```语言名``` 包裹；
4. 重点内容标注 **粗体**/*斜体*；
5. 思考、工具日志仅输出文本，适配前端折叠面板，禁止输出HTML标签；
6. 内容分点拆分，排版整洁适配前端Markdown渲染。

# 行文要求
语言通俗，逻辑完整简洁，无多余废话。
"""


async def run_stream_events(
    task: Task | None,
    user_prompt: str,
    message_history: List[ModelMessage],
) -> AsyncGenerator[AgentStreamEvent | AgentRunResultEvent | TaskNodeResultEvent, None]:
    """
    以流式事件模式运行
    """
    if task:
        match task.type:
            case TaskType.BOOK_FLIGHT:
                from application.tasks.book_flight import run_stream_events

                async for event in run_stream_events(
                    task=task,
                    user_prompt=user_prompt,
                    message_history=message_history,
                ):
                    yield event

    else:
        agent = Agent(
            model=DEEPSEEK_V4_FLASH_MODEL,
            instructions=instruction,
        )
        async with agent.run_stream_events(
            user_prompt=user_prompt,
            message_history=message_history,
        ) as events:
            async for event in events:
                yield event