#  Copyright 2026 zhaoxi826
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

from pydantic_ai import Agent, RunContext
from pydantic import Field
from kilostar.adapter.model_adapter.agent_factory import AgentFactory
from kilostar.core.global_state_machine.model_provider.base_provider import Provider
from kilostar.utils.agent_model import ResponseModel, InputModel, DepsModel
from kilostar.utils.ray_hook import ray_actor_hook

from kilostar.utils.logger import get_logger

logger = get_logger("worker_individual")


class WorkerIndividualResponse(ResponseModel):
    """WorkerIndividualResponse 核心组件类。
    这是一个具体的 Worker 智能体实体类，代表着具备特定人设、领域技能或长文本处理能力的数字员工。它可以被控制器动态拉起，并在安全沙箱内执行复杂的工作流指令与多步骤推理任务。"""

    output: str = Field(..., description="Worker执行任务的输出结果")


class WorkerIndividualDeps(DepsModel):
    """WorkerIndividualDeps 核心组件类。
    这是一个具体的 Worker 智能体实体类，代表着具备特定人设、领域技能或长文本处理能力的数字员工。它可以被控制器动态拉起，并在安全沙箱内执行复杂的工作流指令与多步骤推理任务。"""

    task_event: dict


class WorkerIndividualInput(InputModel):
    """WorkerIndividualInput 核心组件类。
    这是一个具体的 Worker 智能体实体类，代表着具备特定人设、领域技能或长文本处理能力的数字员工。它可以被控制器动态拉起，并在安全沙箱内执行复杂的工作流指令与多步骤推理任务。"""

    task_event: dict


class BaseIndividual:
    """
    Worker Individual 的基类
    """

    def __init__(self, agent_config: dict):
        self.agent_config = agent_config
        self.agent_id = agent_config.get("agent_id")
        self.agent: Agent | None = None

    async def _init_agent(self, agent_name: str, system_prompt: str):
        """完成  agent 模块的启动与依赖初始化。
        在系统引导或服务拉起阶段被调用，负责建立网络连接、分配基础内存资源及注册核心服务组件。
        Args:        agent_name (str): 赋予该实体的人类可读名称或标题字符串，主要用于前端 UI 展示、日志记录或模糊检索。        system_prompt (str): 控制逻辑流向的具体字符串参数，指定了期望的 system prompt 内容。
        Returns:        : 经由当前业务模型加工处理后所输出的具体数据实例或领域模型对象。"""
        from kilostar.utils.get_tool import load_tools_from_list

        global_state_machine = ray_actor_hook(
            "global_state_machine"
        ).global_state_machine
        provider_title = self.agent_config.get(
            "provider_title", "openai"
        )  # default fallback
        model_id = self.agent_config.get("model_id", "gpt-4o")  # default fallback
        tools_list = self.agent_config.get("tools", None)

        provider: Provider = await global_state_machine.get_provider.remote(
            provider_title
        )
        agent_factory = AgentFactory()

        callables = load_tools_from_list(tools_list)

        self.agent = agent_factory.create_agent(
            provider=provider,
            model_id=model_id,
            output_type=WorkerIndividualResponse,
            system_prompt=system_prompt,
            deps_type=WorkerIndividualDeps,
            agent_name=agent_name,
            tools=callables,
        )

        @self.agent.system_prompt
        async def dynamic_prompt(ctx: RunContext[WorkerIndividualDeps]):
            """执行与 dynamic prompt 相关的核心业务流转操作。
            该方法封装了具体的算法策略或状态控制逻辑，确保操作能够在事务上下文中被原子且一致地执行。
            Args:        ctx (RunContext[WorkerIndividualDeps]): 参与 dynamic prompt 逻辑运算或数据构建的上下文依赖对象。
            Returns:        : 经由当前业务模型加工处理后所输出的具体数据实例或领域模型对象。"""
            prompt = system_prompt + "\n\n"
            prompt += f"=== 当前任务上下文 ===\n{ctx.deps.task_event}\n"
            return prompt

    async def run(self, task_event: dict) -> dict:
        """执行与 run 相关的核心业务流转操作。
        该方法封装了具体的算法策略或状态控制逻辑，确保操作能够在事务上下文中被原子且一致地执行。
        Args:        task_event (dict): 由事件总线或工作流引擎分发过来的事件载荷，封装了触发此次调用的上下文快照与任务目标指令。
        Returns:         (dict): 高度聚合的字典结构数据，将多维度的属性特征或统计指标组合后一并返回。"""
        raise NotImplementedError("子类必须实现 run 方法")