# -*- coding: utf-8 -*-
# pylint: disable=C2801, W0611, W0212
from typing import Optional, Any

from agentscope.model import ChatModelBase
from agentscope.formatter import FormatterBase
from agentscope.memory import MemoryBase
from agentscope.tool import ToolResponse
from agentscope.message import (
    Msg,
    TextBlock,
)
from dotenv import load_dotenv

from alias.agent.agents import AliasAgentBase
from alias.agent.tools import AliasToolkit
from alias.agent.utils.constants import WORKER_MAX_ITER
from alias.agent.agents.common_agent_utils import (
    WorkerResponse,
)

load_dotenv()


class ReActWorker(AliasAgentBase):
    def __init__(
        self,
        name: str,
        model: ChatModelBase,
        formatter: FormatterBase,
        memory: MemoryBase,
        toolkit: AliasToolkit,
        sys_prompt: Optional[str] = None,
        max_iters: int = 10,
        state_saving_dir: Optional[str] = None,
        session_service: Any = None,
    ) -> None:
        """Initialize the ReAct agent with the given name, model config name
        and tools.
        """
        super().__init__(
            name=name,
            sys_prompt=sys_prompt,
            model=model,
            formatter=formatter,
            memory=memory,
            toolkit=toolkit,
            max_iters=max_iters,
            session_service=session_service,
            state_saving_dir=state_saving_dir,
        )

        self.max_iters: int = max(self.max_iters, WORKER_MAX_ITER)

    def generate_response(
        self,
        response: str = "",
        task_done: bool = True,
        subtask_progress_summary: str = "",
        generated_files: dict[str, str] = None,
    ) -> ToolResponse:
        """
        Generate a response summarizing the execution progress of the
        given subtask.
        Args:
            response (str):
                The response text (compatible with AgentScope finish function).
            task_done (bool):
                REQUIRED! Whether the subtask was done or not.
            subtask_progress_summary (str):
                REQUIRED! The subtask progress summary.
            generated_files (dict[str, str]):
                REQUIRED! Collect all files generated in the execution process,
                such as the files generated by `write_file` and `edit_file`.
                This field MUST be in dictionary, where the keys are the
                paths of generated files (e.g. '/FULL/PATH/OF/FILE_1.md') and
                the values are short descriptions about the generated files.
        """
        if generated_files is None:
            generated_files = {}

        # If only response is provided,
        # use it as subtask_progress_summary
        if not subtask_progress_summary and response:
            subtask_progress_summary = response

        structure_response = WorkerResponse(
            task_done=task_done,
            subtask_progress_summary=subtask_progress_summary,
            generated_files=generated_files,
        )
        response_msg = Msg(
            self.name,
            content=[
                TextBlock(type="text", text=subtask_progress_summary),
            ],
            role="assistant",
            metadata=structure_response.model_dump(),
        )
        return ToolResponse(
            content=[
                TextBlock(
                    type="text",
                    text="Successfully generated response.",
                ),
            ],
            metadata={
                "success": True,
                "response_msg": response_msg,
            },
            is_last=True,
        )
