#  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, RequestModel, 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):
    """Worker Individual 的输出模型，承载一次任务执行后的结果文本。"""

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


class WorkerIndividualDeps(DepsModel):
    """Worker Individual 的运行期依赖，注入到 pydantic-ai Agent 的 RunContext。"""

    task_event: dict


class WorkerIndividualInput(RequestModel):
    """Worker Individual 的输入模型，承载一次任务事件的入参。"""

    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_config 拉起一个 pydantic-ai Agent 实例。

        从 GlobalStateMachine 取出 Provider，按 agent_config 中的 provider_title
        和 model_id 选择模型，加载工具列表，并把 system_prompt 注册为动态提示词。

        Args:
            agent_name: Agent 的人类可读名称，用于日志与展示。
            system_prompt: 该 Agent 的基础系统提示词，会和 task_event 拼接成动态提示词。
        """
        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]):
            """把基础 system_prompt 与本次 task_event 拼接成最终动态提示词。"""
            prompt = system_prompt + "\n\n"
            prompt += f"=== 当前任务上下文 ===\n{ctx.deps.task_event}\n"
            return prompt

    async def run(self, task_event: dict) -> dict:
        """执行一次任务，需要由子类按自身策略实现。"""
        raise NotImplementedError("子类必须实现 run 方法")