"""Environment management for the Arklex framework.

This module provides functionality for managing the environment, including
worker initialization, tool management, and slot filling integration.
"""

import importlib
import os
import uuid
from collections.abc import Callable
from functools import partial
from typing import Any

from arklex.env.planner.react_planner import DefaultPlanner, ReactPlanner
from arklex.env.tools.tools import Tool
from arklex.env.workers.worker import BaseWorker
from arklex.orchestrator.NLU.core.slot import SlotFiller
from arklex.orchestrator.NLU.services.api_service import APIClientService
from arklex.orchestrator.NLU.services.model_service import (
    DummyModelService,
    ModelService,
)
from arklex.utils.graph_state import MessageState, NodeInfo, Params
from arklex.utils.logging_utils import LogContext

log_context = LogContext(__name__)


class BaseResourceInitializer:
    """Abstract base class for resource initialization.

    This class defines the interface for initializing tools and workers in the environment.
    Concrete implementations must provide methods for tool and worker initialization.
    """

    @staticmethod
    def init_tools(tools: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
        """Initialize tools from configuration.

        Args:
            tools: List of tool configurations

        Returns:
            Dictionary mapping tool IDs to their configurations

        Raises:
            NotImplementedError: Must be implemented by subclasses
        """
        raise NotImplementedError

    @staticmethod
    def init_workers(workers: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
        """Initialize workers from configuration.

        Args:
            workers: List of worker configurations

        Returns:
            Dictionary mapping worker IDs to their configurations

        Raises:
            NotImplementedError: Must be implemented by subclasses
        """
        raise NotImplementedError


class DefaultResourceInitializer(BaseResourceInitializer):
    """Default implementation of resource initialization.

    This class provides a default implementation for initializing tools and workers
    in the environment.
    """

    @staticmethod
    def init_tools(tools: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
        """Initialize tools from configuration.

        Args:
            tools: List of tool configurations

        Returns:
            Dictionary mapping tool IDs to their configurations
        """
        tool_registry: dict[str, dict[str, Any]] = {}
        for tool in tools:
            tool_id: str = tool["id"]
            name: str = tool["name"]
            path: str = tool["path"]
            try:
                filepath: str = os.path.join("arklex.env.tools", path)
                module_name: str = filepath.replace(os.sep, ".").replace(".py", "")
                module = importlib.import_module(module_name)
                func: Callable = getattr(module, name)
                tool_instance: Tool = func()
                tool_registry[tool_id] = {
                    "name": f"{path.replace('/', '-')}-{name}",
                    "description": tool_instance.description,
                    "execute": func,
                    "fixed_args": tool.get("fixed_args", {}),
                }
            except Exception as e:
                log_context.error(f"Tool {name} is not registered, error: {e}")
        return tool_registry

    @staticmethod
    def init_workers(workers: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
        """Initialize workers from configuration.

        Args:
            workers: List of worker configurations

        Returns:
            Dictionary mapping worker IDs to their configurations
        """
        worker_registry: dict[str, dict[str, Any]] = {}
        for worker in workers:
            worker_id: str = worker["id"]
            name: str = worker["name"]
            path: str = worker["path"]
            try:
                filepath: str = os.path.join("arklex.env.workers", path)
                module_name: str = filepath.replace(os.sep, ".").rstrip(".py")
                module = importlib.import_module(module_name)
                func: Callable = getattr(module, name)
                worker_registry[worker_id] = {
                    "name": name,
                    "description": func.description,
                    "execute": partial(func, **worker.get("fixed_args", {})),
                }
            except Exception as e:
                log_context.error(f"Worker {name} is not registered, error: {e}")
        return worker_registry


class Environment:
    """Environment management for workers and tools.

    This class manages the environment for workers and tools, including
    initialization, state management, and slot filling integration.
    """

    def __init__(
        self,
        tools: list[dict[str, Any]],
        workers: list[dict[str, Any]],
        slotsfillapi: str = "",
        resource_initializer: BaseResourceInitializer | None = None,
        planner_enabled: bool = False,
        model_service: ModelService | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize the environment.

        Args:
            tools: List of tools to initialize
            workers: List of workers to initialize
            slotsfillapi: API endpoint for slot filling
            resource_initializer: Resource initializer instance
            planner_enabled: Whether planning is enabled
            model_service: Model service for intent detection and slot filling
        """
        # Accept slot_fill_api as an alias for slotsfillapi for compatibility with tests
        if "slot_fill_api" in kwargs and not slotsfillapi:
            slotsfillapi = kwargs["slot_fill_api"]
        if resource_initializer is None:
            resource_initializer = DefaultResourceInitializer()
        self.tools: dict[str, dict[str, Any]] = resource_initializer.init_tools(tools)
        self.workers: dict[str, dict[str, Any]] = resource_initializer.init_workers(
            workers
        )
        self.name2id: dict[str, str] = {
            resource["name"]: id
            for id, resource in {**self.tools, **self.workers}.items()
        }
        self.id2name: dict[str, str] = {
            id: resource["name"]
            for id, resource in {**self.tools, **self.workers}.items()
        }
        self.model_service = model_service or DummyModelService(
            {
                "model_name": "dummy",
                "api_key": "dummy",
                "endpoint": "http://dummy",
                "model_type_or_path": "dummy-path",
                "llm_provider": "dummy",
            }
        )
        self.slotfillapi: SlotFiller = self.initialize_slotfillapi(slotsfillapi)
        if planner_enabled:
            self.planner: ReactPlanner | DefaultPlanner = ReactPlanner(
                tools_map=self.tools, workers_map=self.workers, name2id=self.name2id
            )
        else:
            self.planner: ReactPlanner | DefaultPlanner = DefaultPlanner(
                tools_map=self.tools, workers_map=self.workers, name2id=self.name2id
            )

    def initialize_slotfillapi(self, slotsfillapi: str) -> SlotFiller:
        """Initialize the slot filling API.

        Args:
            slotsfillapi: API endpoint for slot filling. If not a string or empty,
                         falls back to local model-based slot filling.

        Returns:
            SlotFiller: Initialized slot filler instance, either API-based or local model-based.
        """
        if isinstance(slotsfillapi, str) and slotsfillapi:
            api_service = APIClientService(base_url=slotsfillapi)
            return SlotFiller(model_service=self.model_service, api_service=api_service)
        else:
            return SlotFiller(model_service=self.model_service)

    def step(
        self, id: str, message_state: MessageState, params: Params, node_info: NodeInfo
    ) -> tuple[MessageState, Params]:
        """Execute a step in the environment.

        Args:
            id: Resource ID to execute
            message_state: Current message state
            params: Current parameters
            node_info: Information about the current node

        Returns:
            Tuple containing updated message state and parameters
        """
        response_state: MessageState
        if id in self.tools:
            log_context.info(f"{self.tools[id]['name']} tool selected")
            tool: Tool = self.tools[id]["execute"]()
            tool.init_slotfiller(self.slotfillapi)
            tool.load_slots(getattr(node_info, 'attributes', {}).get("slots", []))
            combined_args: dict[str, Any] = {
                **self.tools[id]["fixed_args"],
                **(node_info.additional_args or {}),
            }
            response_state = tool.execute(message_state, **combined_args)
            params.memory.function_calling_trajectory = (
                response_state.function_calling_trajectory
            )
            params.taskgraph.dialog_states = response_state.slots
            params.taskgraph.node_status[params.taskgraph.curr_node] = (
                response_state.status
            )

        elif id in self.workers:
            log_context.info(f"{self.workers[id]['name']} worker selected")
            worker: BaseWorker = self.workers[id]["execute"]()
            if hasattr(worker, "init_slotfilling"):
                worker.init_slotfilling(self.slotfillapi)
            response_state = worker.execute(message_state, **node_info.additional_args)
            call_id: str = str(uuid.uuid4())
            params.memory.function_calling_trajectory.append(
                {
                    "content": None,
                    "role": "assistant",
                    "tool_calls": [
                        {
                            "function": {"arguments": "{}", "name": self.id2name[id]},
                            "id": call_id,
                            "type": "function",
                        }
                    ],
                    "function_call": None,
                }
            )
            params.memory.function_calling_trajectory.append(
                {
                    "role": "tool",
                    "tool_call_id": call_id,
                    "name": self.id2name[id],
                    "content": response_state.response
                    if response_state.response
                    else response_state.message_flow,
                }
            )
            params.taskgraph.node_status[params.taskgraph.curr_node] = (
                response_state.status
            )
        else:
            log_context.info("planner selected")
            action: str
            response_state: MessageState
            msg_history: list[dict[str, Any]]
            action, response_state, msg_history = self.planner.execute(
                message_state, params.memory.function_calling_trajectory
            )

        log_context.info(f"Response state from {id}: {response_state}")
        return response_state, params

    def register_tool(self, name: str, tool: Tool) -> None:
        """Register a tool in the environment.

        Args:
            name: Name of the tool
            tool: Tool instance

        Raises:
            EnvironmentError: If tool registration fails
        """
        try:
            self.tools[name] = tool
            log_context.info(f"{self.tools[name]['name']} tool selected")
        except Exception as e:
            log_context.error(f"Tool {name} is not registered, error: {e}")
