import logging
from typing import List, Type, Tuple, Dict, Any, Optional
import importlib

from pydantic import BaseModel, Field, PrivateAttr, model_validator

from moatless.actions.action import Action
from moatless.actions.model import ActionArguments, Observation
from moatless.actions.reject import RejectArgs
from moatless.completion.completion import (
    CompletionModel,
)
from moatless.completion.model import Message, AssistantMessage, UserMessage, Completion
from moatless.index.code_index import CodeIndex
from moatless.node import Node
from moatless.repository.repository import Repository

logger = logging.getLogger(__name__)


class ActionAgent(BaseModel):
    system_prompt: Optional[str] = Field(
        None, description="System prompt to be used for generating completions"
    )
    actions: List[Action] = Field(default_factory=list)

    _completion: CompletionModel = PrivateAttr()
    _action_map: dict[Type[ActionArguments], Action] = PrivateAttr(default_factory=dict)

    def __init__(
        self,
        system_prompt: str | None = None,
        actions: List[Action] | None = None,
        completion: CompletionModel | None = None,
    ):
        actions = actions or []
        actions_map = {action.args_schema: action for action in actions}
        super().__init__(actions=actions, system_prompt=system_prompt)
        self._completion = completion
        self._action_map = actions_map

    @model_validator(mode="after")
    def verify_system_prompt(self) -> "ActionAgent":
        if self.system_prompt == "":
            self.system_prompt = None
        return self

    def run(self, node: Node):
        self._generate_action(node)
        self._execute_action(node)

        logger.info(
            f"Node{node.node_id}: Executed action: {node.action.name}. "
            f"Terminal: {node.observation.terminal}. "
            f"Output: {node.observation.message}"
        )

    def _generate_action(self, node: Node):
        """
        Generate an action
        """

        if node.action:
            logger.info(f"Node{node.node_id}: Action already generated. Skipping.")
            return

        completion_response = None
        try:
            possible_actions = self._determine_possible_actions(node)
            node.possible_actions = [action.name for action in possible_actions]

            system_prompt = self._create_system_prompt(possible_actions)
            messages = self._create_messages(node)

            action_args, completion_response = self._generate_action_args(
                system_prompt, messages, possible_actions
            )

            node.action = action_args
            node.completions["build_action"] = completion_response

            duplicate_node = node.find_duplicate()
            if duplicate_node:
                logger.info(
                    f"Node{node.node_id} is a duplicate to Node{duplicate_node.node_id}. Skipping execution."
                )
                node.is_duplicate = True
                return

        except Exception as e:
            logger.exception(f"Node{node.node_id}: Error generating action. Set action to Reject")
            error_message = f"Failed to generate action: {str(e)}"
            node.action = RejectArgs(
                rejection_reason=error_message,
                scratch_pad=f"An error occurred during action generation: {error_message}",
            )

    def _execute_action(self, node: Node):
        if node.observation:
            logger.info(f"Node{node.node_id}: Observation already generated. Skipping.")
            return

        action = self._action_map.get(type(node.action))
        if action:
            node.observation = action.execute(node.action, node.file_context)

            if node.observation.execution_completion:
                node.completions["execute_action"] = (
                    node.observation.execution_completion
                )
        else:
            logger.error(
                f"Node{node.node_id}: Action {node.action} not found in action map. "
                f"Available actions: {self._action_map.keys()}"
            )
            raise Exception(f"Action {node.action} not found in action map.")

    def _create_system_prompt(self, possible_actions: List[Action]) -> str:
        return self.system_prompt

    def _create_messages(self, node: Node) -> list[Message]:
        messages: list[Message] = []

        last_node = None
        previous_nodes = node.get_trajectory()[:-1]
        for previous_node in previous_nodes:
            if previous_node.action:
                tool_call = previous_node.action.to_tool_call()
                messages.append(AssistantMessage(tool_call=tool_call))

            content = ""
            if previous_node.message:
                # TODO: Don't hardcode the issue tag...
                content = f"<issue>\n{previous_node.message}\n</issue>"

            if previous_node.observation:
                content += previous_node.observation.message

            if not content:
                logger.warning(
                    f"Node{previous_node.node_id}: No content to add to messages"
                )

            messages.append(UserMessage(content=content))

            last_node = previous_node

        if last_node.observation and last_node.observation.extra:
            messages[-1].content += "\n" + last_node.observation.extra

        # TODO: Only add the updated file context per node
        if node.file_context:
            if node.file_context.is_empty():
                file_context_str = "No files added to file context yet."
            else:
                file_context_str = node.file_context.create_prompt(
                    show_span_ids=False,
                    show_line_numbers=True,
                    exclude_comments=False,
                    show_outcommented_code=True,
                    outcomment_code_comment="... rest of the code",
                )

            messages[-1].content = (
                f"# Current file context:\n\n<file_context>\n{file_context_str}\n</file_context>\n\n"
                + messages[-1].content
            )

        if node.feedback:
            logger.info(f"Node{node.node_id}: Feedback provided: {node.feedback}")
            messages[-1].content += f"\n\n{node.feedback}"

        return messages

    def _generate_action_args(
        self, system_prompt: str, messages: List[Message], actions: List[Action]
    ) -> Tuple[ActionArguments, Completion]:
        try:
            action_args = []
            for action in actions:
                if not isinstance(action, Action):
                    raise TypeError(
                        f"Invalid action type: {type(action)}. Expected Action subclass."
                    )
                if not hasattr(action, "args_schema"):
                    raise AttributeError(
                        f"Action {action.__class__.__name__} is missing args_schema attribute"
                    )
                action_args.append(action.args_schema)

            return self._completion.create_completion(
                messages, system_prompt=system_prompt, actions=action_args
            )
        except Exception as e:
            logger.exception(f"Error in _generate_action_args: {str(e)}")
            problematic_actions = [
                f"{action.__class__.__name__} (type: {type(action)})"
                for action in actions
                if not isinstance(action, Action) or not hasattr(action, "args_schema")
            ]
            if problematic_actions:
                error_message = f"The following actions are invalid or missing args_schema attribute: {', '.join(problematic_actions)}"
            else:
                error_message = (
                    "Unknown error occurred while generating action arguments"
                )
            raise Exception(error_message) from e

    def _determine_possible_actions(self, node: Node) -> List[Action]:
        actions = self.actions
        logger.debug(
            f"Possible actions for Node{node.node_id}: {[action.__class__.__name__ for action in actions]}"
        )
        return actions

    def model_dump(self, **kwargs) -> Dict[str, Any]:
        dump = super().model_dump(**kwargs)
        dump["completion"] = self._completion.model_dump(**kwargs)
        dump["actions"] = []
        dump["agent_class"] = f"{self.__class__.__module__}.{self.__class__.__name__}"
        for action in self.actions:
            action_dump = action.model_dump(**kwargs)
            action_dump["action_class"] = (
                f"{action.__class__.__module__}.{action.__class__.__name__}"
            )
            dump["actions"].append(action_dump)
        return dump

    @classmethod
    def model_validate(
        cls,
        obj: Any,
        repository: Repository = None,
        runtime: Any = None,
        code_index: CodeIndex = None,
    ) -> "ActionAgent":
        if isinstance(obj, dict):
            obj = obj.copy()
            completion_data = obj.pop("completion", None)
            agent_class_path = obj.pop("agent_class", None)

            if completion_data:
                completion = CompletionModel.model_validate(completion_data)
            else:
                completion = None

            if repository:
                actions = [
                    Action.from_dict(
                        action_data,
                        repository=repository,
                        runtime=runtime,
                        code_index=code_index,
                    )
                    for action_data in obj.get("actions", [])
                ]
            else:
                logger.debug(f"No repository provided, skip initiating actions")
                actions = []

            if agent_class_path:
                module_name, class_name = agent_class_path.rsplit(".", 1)
                module = importlib.import_module(module_name)
                agent_class = getattr(module, class_name)
                instance = agent_class(actions=actions, completion=completion)
            else:
                instance = cls(actions=actions, completion=completion)

            instance._action_map = {action.args_schema: action for action in actions}
            return instance

        return super().model_validate(obj)

    # Add a property for completion to allow access while maintaining it as a private attribute
    @property
    def completion(self) -> CompletionModel:
        return self._completion

    @completion.setter
    def completion(self, value: CompletionModel):
        self._completion = value
