import json

from superagi.agent.task_queue import TaskQueue
from superagi.agent.agent_message_builder import AgentLlmMessageBuilder
from superagi.agent.agent_prompt_builder import AgentPromptBuilder
from superagi.agent.output_handler import ToolOutputHandler
from superagi.agent.output_parser import AgentSchemaToolOutputParser
from superagi.agent.queue_step_handler import QueueStepHandler
from superagi.agent.tool_builder import ToolBuilder
from superagi.helper.error_handler import ErrorHandler
from superagi.helper.prompt_reader import PromptReader
from superagi.helper.token_counter import TokenCounter
from superagi.lib.logger import logger
from superagi.models.agent import Agent
from superagi.models.agent_config import AgentConfiguration
from superagi.models.agent_execution import AgentExecution
from superagi.models.agent_execution_config import AgentExecutionConfiguration
from superagi.models.agent_execution_feed import AgentExecutionFeed
from superagi.models.agent_execution_permission import AgentExecutionPermission
from superagi.models.tool import Tool
from superagi.models.toolkit import Toolkit
from superagi.models.workflows.agent_workflow_step import AgentWorkflowStep
from superagi.models.workflows.agent_workflow_step_tool import AgentWorkflowStepTool
from superagi.resource_manager.resource_summary import ResourceSummarizer
from superagi.tools.base_tool import BaseTool
from sqlalchemy import and_

class AgentToolStepHandler:
    """Handles the tools steps in the agent workflow"""
    def __init__(self, session, llm, agent_id: int, agent_execution_id: int, memory=None):
        self.session = session
        self.llm = llm
        self.agent_execution_id = agent_execution_id
        self.agent_id = agent_id
        self.memory = memory
        self.task_queue = TaskQueue(str(self.agent_execution_id))
        self.organisation = Agent.find_org_by_agent_id(self.session, self.agent_id)

    def execute_step(self):
        execution = AgentExecution.get_agent_execution_from_id(self.session, self.agent_execution_id)
        workflow_step = AgentWorkflowStep.find_by_id(self.session, execution.current_agent_step_id)
        step_tool = AgentWorkflowStepTool.find_by_id(self.session, workflow_step.action_reference_id)
        agent_config = Agent.fetch_configuration(self.session, self.agent_id)
        agent_execution_config = AgentExecutionConfiguration.fetch_configuration(self.session, self.agent_execution_id)
        # print(agent_execution_config)

        if not self._handle_wait_for_permission(execution, workflow_step):
            return

        if step_tool.tool_name == "TASK_QUEUE":
            step_response = QueueStepHandler(self.session, self.llm, self.agent_id, self.agent_execution_id).execute_step()
            next_step = AgentWorkflowStep.fetch_next_step(self.session, workflow_step.id, step_response)
            self._handle_next_step(next_step)
            return

        if step_tool.tool_name == "WAIT_FOR_PERMISSION":
            self._create_permission_request(execution, step_tool)
            return

        assistant_reply = self._process_input_instruction(agent_config, agent_execution_config, step_tool,
                                                          workflow_step)
        tool_obj = self._build_tool_obj(agent_config, agent_execution_config, step_tool.tool_name)
        tool_output_handler = ToolOutputHandler(self.agent_execution_id, agent_config, [tool_obj],self.memory,
                                                output_parser=AgentSchemaToolOutputParser())
        final_response = tool_output_handler.handle(self.session, assistant_reply)
        step_response = "default"
        if step_tool.output_instruction:
            step_response = self._process_output_instruction(final_response.result, step_tool, workflow_step)

        next_step = AgentWorkflowStep.fetch_next_step(self.session, workflow_step.id, step_response)
        self._handle_next_step(next_step)
        self.session.flush()

    def _create_permission_request(self, execution, step_tool: AgentWorkflowStepTool):
        new_agent_execution_permission = AgentExecutionPermission(
            agent_execution_id=self.agent_execution_id,
            status="PENDING",
            agent_id=self.agent_id,
            tool_name="WAIT_FOR_PERMISSION",
            question=step_tool.input_instruction,
            assistant_reply="")
        self.session.add(new_agent_execution_permission)
        self.session.commit()
        self.session.flush()
        execution.permission_id = new_agent_execution_permission.id
        execution.status = "WAITING_FOR_PERMISSION"
        self.session.commit()

    def _handle_next_step(self, next_step):
        if str(next_step) == "COMPLETE":
            agent_execution = AgentExecution.get_agent_execution_from_id(self.session, self.agent_execution_id)
            agent_execution.current_agent_step_id = -1
            agent_execution.status = "COMPLETED"
        else:
            AgentExecution.assign_next_step_id(self.session, self.agent_execution_id, next_step.id)
        self.session.commit()

    def _process_input_instruction(self, agent_config, agent_execution_config, step_tool, workflow_step):
        tool_obj = self._build_tool_obj(agent_config, agent_execution_config, step_tool.tool_name)
        prompt = self._build_tool_input_prompt(step_tool, tool_obj, agent_execution_config)
        logger.info("Prompt: ", prompt)
        agent_feeds = AgentExecutionFeed.fetch_agent_execution_feeds(self.session, self.agent_execution_id)
        messages = AgentLlmMessageBuilder(self.session, self.llm, self.llm.get_model(), self.agent_id, self.agent_execution_id) \
            .build_agent_messages(prompt, agent_feeds, history_enabled=step_tool.history_enabled,
                                  completion_prompt=step_tool.completion_prompt)
        # print(messages)
        current_tokens = TokenCounter.count_message_tokens(messages, self.llm.get_model())
        response = self.llm.chat_completion(messages, TokenCounter(session=self.session, organisation_id=self.organisation.id).token_limit(self.llm.get_model()) - current_tokens)

        if 'error' in response and response['message'] is not None:
            ErrorHandler.handle_openai_errors(self.session, self.agent_id, self.agent_execution_id, response['message'])
        # ModelsHelper(session=self.session, organisation_id=organisation.id).create_call_log(execution.name,agent_config['agent_id'],response['response'].usage.total_tokens,json.loads(response['content'])['tool']['name'],agent_config['model'])
        if 'content' not in response or response['content'] is None:
            raise RuntimeError(f"Failed to get response from llm")
        total_tokens = current_tokens + TokenCounter.count_message_tokens(response, self.llm.get_model())
        AgentExecution.update_tokens(self.session, self.agent_execution_id, total_tokens)
        assistant_reply = response['content']
        return assistant_reply

    def _build_tool_obj(self, agent_config, agent_execution_config, tool_name: str):
        model_api_key = AgentConfiguration.get_model_api_key(self.session, self.agent_id, agent_config["model"])['api_key']
        tool_builder = ToolBuilder(self.session, self.agent_id, self.agent_execution_id)
        resource_summary = ""
        if tool_name == "QueryResourceTool":
            resource_summary = ResourceSummarizer(session=self.session,
                                                  agent_id=self.agent_id,
                                                  model=agent_config["model"]).fetch_or_create_agent_resource_summary(
                default_summary=agent_config.get("resource_summary"))

        organisation = Agent.find_org_by_agent_id(self.session, self.agent_id)
        tool = self.session.query(Tool).join(Toolkit, and_(Tool.toolkit_id == Toolkit.id, Toolkit.organisation_id == organisation.id, Tool.name == tool_name)).first()
        tool_obj = tool_builder.build_tool(tool)
        tool_obj = tool_builder.set_default_params_tool(tool_obj, agent_config, agent_execution_config, model_api_key,
                                                        resource_summary,self.memory)
        return tool_obj

    def _process_output_instruction(self, final_response: str, step_tool: AgentWorkflowStepTool,
                                    workflow_step: AgentWorkflowStep):
        prompt = self._build_tool_output_prompt(step_tool, final_response, workflow_step)
        messages = [{"role": "system", "content": prompt}]
        current_tokens = TokenCounter.count_message_tokens(messages, self.llm.get_model())
        response = self.llm.chat_completion(messages,
                                            TokenCounter(session=self.session, organisation_id=self.organisation.id).token_limit(self.llm.get_model()) - current_tokens)

        if 'error' in response and response['message'] is not None:
            ErrorHandler.handle_openai_errors(self.session, self.agent_id, self.agent_execution_id, response['message'])
            
        if 'content' not in response or response['content'] is None:
            raise RuntimeError(f"ToolWorkflowStepHandler: Failed to get output response from llm")
        total_tokens = current_tokens + TokenCounter.count_message_tokens(response, self.llm.get_model())
        AgentExecution.update_tokens(self.session, self.agent_execution_id, total_tokens)
        step_response = response['content']
        step_response = step_response.replace("'", "").replace("\"", "")
        return step_response

    def _build_tool_input_prompt(self, step_tool: AgentWorkflowStepTool, tool: BaseTool, agent_execution_config: dict):
        super_agi_prompt = PromptReader.read_agent_prompt(__file__, "agent_tool_input.txt")
        super_agi_prompt = super_agi_prompt.replace("{goals}", AgentPromptBuilder.add_list_items_to_string(
            agent_execution_config["goal"]))
        super_agi_prompt = super_agi_prompt.replace("{tool_name}", step_tool.tool_name)
        super_agi_prompt = super_agi_prompt.replace("{instruction}", step_tool.input_instruction)

        tool_schema = f"\"{tool.name}\": {tool.description}, args json schema: {json.dumps(tool.args)}"
        super_agi_prompt = super_agi_prompt.replace("{tool_schema}", tool_schema)
        return super_agi_prompt

    def _get_step_responses(self, workflow_step: AgentWorkflowStep):
        return [step["step_response"] for step in workflow_step.next_steps]

    def _build_tool_output_prompt(self, step_tool: AgentWorkflowStepTool, tool_output: str,
                                  workflow_step: AgentWorkflowStep):
        super_agi_prompt = PromptReader.read_agent_prompt(__file__, "agent_tool_output.txt")
        super_agi_prompt = super_agi_prompt.replace("{tool_output}", tool_output)
        super_agi_prompt = super_agi_prompt.replace("{tool_name}", step_tool.tool_name)
        super_agi_prompt = super_agi_prompt.replace("{instruction}", step_tool.output_instruction)

        step_responses = self._get_step_responses(workflow_step)
        if "default" in step_responses:
            step_responses.remove("default")
        super_agi_prompt = super_agi_prompt.replace("{output_options}", str(step_responses))
        return super_agi_prompt

    def _handle_wait_for_permission(self, agent_execution, workflow_step: AgentWorkflowStep):
        """
        Handles the wait for permission when the agent execution is waiting for permission.

        Args:
            agent_execution (AgentExecution): The agent execution.
            workflow_step (AgentWorkflowStep): The workflow step.

        Raises:
            Returns permission success or failure
        """
        if agent_execution.status != "WAITING_FOR_PERMISSION":
            return True
        agent_execution_permission = self.session.query(AgentExecutionPermission).filter(
            AgentExecutionPermission.id == agent_execution.permission_id).first()
        if agent_execution_permission.status == "PENDING":
            logger.error("handle_wait_for_permission: Permission is still pending")
            return False
        if agent_execution_permission.status == "APPROVED":
            next_step = AgentWorkflowStep.fetch_next_step(self.session, workflow_step.id, "YES")
        else:
            next_step = AgentWorkflowStep.fetch_next_step(self.session, workflow_step.id, "NO")
            result = f"{' User has given the following feedback : ' + agent_execution_permission.user_feedback if agent_execution_permission.user_feedback else ''}"


            agent_execution_feed = AgentExecutionFeed(agent_execution_id=agent_execution_permission.agent_execution_id,
                                                      agent_id=agent_execution_permission.agent_id,
                                                      feed=result, role="user",
                                                      feed_group_id=agent_execution.current_feed_group_id)
            self.session.add(agent_execution_feed)

        agent_execution.status = "RUNNING"
        agent_execution.permission_id = -1
        self.session.commit()
        self._handle_next_step(next_step)
        self.session.commit()
        return False
