#  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 pretor.adapter.model_adapter.agent_factory import AgentFactory
from pretor.core.global_state_machine.model_provider.base_provider import Provider
from pretor.utils.agent_model import ResponseModel, InputModel, DepsModel
from pretor.utils.ray_hook import ray_actor_hook

from pretor.utils.logger import get_logger
logger = get_logger('worker_individual')

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

class WorkerIndividualDeps(DepsModel):
    task_event: dict

class WorkerIndividualInput(InputModel):
    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):
        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

        provider: Provider = await global_state_machine.get_provider.remote( provider_title)
        agent_factory = AgentFactory()
        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
        )

        @self.agent.system_prompt
        async def dynamic_prompt(ctx: RunContext[WorkerIndividualDeps]):
            prompt = system_prompt + "\n\n"
            prompt += (
                f"=== 当前任务上下文 ===\n"
                f"{ctx.deps.task_event}\n"
            )
            return prompt

    async def run(self, task_event: dict) -> dict:
        raise NotImplementedError("子类必须实现 run 方法")