import asyncio
import json
import logging
import uuid
from json import JSONDecodeError
from typing import Literal, Any

from superagentx.browser_engine import BrowserEngine
from superagentx.task_engine import TaskEngine
from superagentx.constants import SEQUENCE, PARALLEL
from superagentx.engine import Engine
from superagentx.exceptions import StopSuperAgentX
from superagentx.llm import LLMClient, ChatCompletionParams
from superagentx.prompt import PromptTemplate
from superagentx.result import GoalResult
from superagentx.utils.helper import iter_to_aiter, StatusCallback, _maybe_await

logger = logging.getLogger(__name__)

_GOAL_PROMPT_TEMPLATE = """Review the given output context and make sure

the following goal is achieved.

Goal: {goal}

Query_Instruction: {query_instruction}

Output_Context : {output_context}

Feedback: {feedback}

Output_Format: {output_format}

Follow the instructions step-by-step carefully and act upon.

Review the Output_Context based on the given Goal with Query_Instruction and set the result in the below mentioned result.

Answer should be based on the given output context. Do not try answer by your own.

Make sure generate the result based on the given output format if provided. 

{result_format}

Always generate the JSON output. Don't include any command lines.
"""

ENGINE_RESULT_FORMAT = """
{{
    reason: Set the reason for result,
    result: Set this based on given output format if output format given. Otherwise set the result as it is.,
    is_goal_satisfied: 'True' if result satisfied based on the given goal. Otherwise set as 'False'. Set only 'True' or 'False' boolean.
}}
"""


class Agent:

    def __init__(
            self,
            *,
            goal: str | None = None,
            role: str | None = None,
            llm: LLMClient | None = None,
            prompt_template: PromptTemplate | None = None,
            agent_id: str | None = None,
            name: str | None = None,
            description: str | None = None,
            engines: list[Engine | BrowserEngine | TaskEngine | list[Engine | BrowserEngine | TaskEngine]] | None = None,
            output_format: str | None = None,
            max_retry: int = 5,
            return_engine_result: bool = False
    ):
        """
        Initializes a new instance of the Agent class.

        This constructor sets up an engine with specific goals and roles, allowing
        it to interact with a large language model (LLM) and utilize a defined
        prompt template. The engine can be configured with various parameters to
        tailor its behavior to specific tasks and workflows.

        Args:
            goal: The primary objective or goal that the engine is designed to achieve.
            role: The role or function that the engine will assume in its operations.
            llm: Interface for communicating with the large language model (LLM).
            prompt_template: Defines the structure and format of prompts sent to the LLM using `PromptTemplate`.
            agent_id: A unique identifier for the engine. If not provided, a new UUID
                will be generated by default. Useful for tracking or referencing
                the engine in multi-engine environments.
            name: An optional name for the engine, providing a more friendly reference for display or logging purposes.
            description: An optional description that provides additional context or details about the engine's
                purpose and capabilities.
            engines: A list of engines (or lists of engines) that the engine can utilize.
                This allows for flexibility in processing and task execution based on different capabilities
                or configurations.
            output_format: Specifies the desired format for the engine's output. This can dictate how results are
                structured and presented.
            max_retry: The maximum number of retry attempts for operations that may fail.
                Default is set to 5. This is particularly useful in scenarios where transient errors may occur,
                ensuring robust execution.
        """
        self.role = role
        self.goal = goal
        self.llm = llm
        self.prompt_template = prompt_template
        self.agent_id = agent_id or uuid.uuid4().hex
        self.name = name or f'{self.__str__()}-{self.agent_id}'
        self.description = description
        self.engines: list[Engine | BrowserEngine | TaskEngine | list[Engine | BrowserEngine | TaskEngine]] = engines or []
        self.output_format = output_format
        self.max_retry = max_retry if max_retry >= 1 else 1
        self.return_engine_result = return_engine_result
        self.engine_result_format = ENGINE_RESULT_FORMAT
        if self.return_engine_result:
            self.engine_result_format = """{{ reason: Set the reason for result, is_goal_satisfied: 'True' if result 
            satisfied based on the given goal. Otherwise set as 'False'. Set only 'True' or 'False' boolean. }}"""

    def __str__(self):
        return "Agent"

    def __repr__(self):
        return f"<{self.__str__()}>"

    async def add(
            self,
            *engines: Engine | BrowserEngine | TaskEngine,
            execute_type: Literal['SEQUENCE', 'PARALLEL'] = 'SEQUENCE'
    ) -> None:
        """
        Adds one or more Engine instances to the current context for processing.

        This method allows the user to include multiple engines that will be used
        for execution based on the specified execution type. The `execute_type`
        parameter determines how the engines will be run: either in a sequence,
        where each engine runs one after the other, or in parallel, where all
        specified engines run concurrently.

        Args:
            engines: One or more Engine instances to be added to the context.
                This allows for flexibility in processing and task execution based on different capabilities
                or configurations.
            execute_type: The method of execution for the added engines.
                - 'SEQUENCE': Engines are executed one after another,
                  waiting for each to complete before starting the next.
                - 'PARALLEL': All engines are executed concurrently, allowing for
                  simultaneous processing.
                Default is 'SEQUENCE'.

        Returns:
            None
        """
        if execute_type == SEQUENCE:
            self.engines += engines
            logger.debug(f'Engine(s) added as {SEQUENCE} : {",".join([str(_engine) for _engine in engines])}')
        else:
            self.engines.append(list(engines))
            logger.debug(f'Engines added as {PARALLEL} : {",".join([str(_engine) for _engine in engines])}')

    async def _verify_goal(
            self,
            *,
            query_instruction: str,
            results: list[Any],
            old_memory: str | None = None,
            pipe_id: str | None = None,
            status_callback: StatusCallback | None = None

    ) -> GoalResult | None:
        if old_memory:
            results = f"output_context:\n{old_memory}\n\n{results}"
            logger.debug(f'Updated Output Context with old memory : {results}')

        prompt_message = await self.prompt_template.get_messages(
            input_prompt=_GOAL_PROMPT_TEMPLATE,
            goal=self.goal,
            query_instruction=query_instruction,
            output_context=results,
            feedback="",
            output_format=self.output_format or "",
            result_format=self.engine_result_format
        )
        chat_completion_params = ChatCompletionParams(
            messages=prompt_message
        )
        messages = await self.llm.achat_completion(
            chat_completion_params=chat_completion_params
        )
        # Callback: agent execution started
        if status_callback:
            await _maybe_await(status_callback(
                event="agent_verify_goal_usage",
                pipe_id=pipe_id,
                agent_id=self.agent_id,
                agent=self.name,
                query=query_instruction,
                messages=messages
            ))

        if messages and messages.choices:
            for choice in messages.choices:
                if choice and choice.message:
                    _res = choice.message.content or ''
                    _res = _res.replace('```json', '').replace('```', '')
                    try:
                        __res = json.loads(_res)
                        return GoalResult(
                            name=self.name,
                            agent_id=self.agent_id,
                            **__res
                        )
                    except JSONDecodeError as ex:
                        _msg = f'Cannot verify goal!\n{ex}'
                        logger.warning(_msg)
                        return GoalResult(
                            name=self.name,
                            agent_id=self.agent_id,
                            content=_res,
                            error=_msg
                        )
        else:
            return GoalResult(
                name=self.name,
                agent_id=self.agent_id,
                error='No results found!',
                is_goal_satisfied=False
            )

    async def _execute(
            self,
            query_instruction: str,
            pre_result: str | None = None,
            old_memory: list[dict] | None = None,
            verify_goal: bool = True,
            pipe_id: str | None = None,
            conversation_id: str | None = None,
            status_callback: StatusCallback | None = None

    ) -> GoalResult:
        results = []
        params = {
            "input_prompt": query_instruction,
            "pre_result": pre_result,
            "old_memory": old_memory,
            "pipe_id": pipe_id,
            "agent_id": self.agent_id,
            "agent_name": self.name,
            "status_callback": status_callback
        }
        if conversation_id:
            params["conversation_id"] = conversation_id
        async for _engines in iter_to_aiter(self.engines):
            if isinstance(_engines, list):
                logger.debug(f'Engine(s) are executing : {",".join([str(_engine) for _engine in _engines])}')
                _res = await asyncio.gather(
                    *[
                        _engine.start(**params)
                        async for _engine in iter_to_aiter(_engines)
                    ]
                )
                logger.debug(f'Engine(s) results : {_res}')
            else:
                logger.debug(f'Engine is executing : {_engines}')
                _res = await _engines.start(**params)
                logger.debug(f'Engine result : {_res}')
            results.append(_res)

        logger.debug(f'Verifying agent goal `{verify_goal}`')
        if verify_goal:
            final_result = await self._verify_goal(
                results=results,
                query_instruction=query_instruction,
                old_memory=old_memory,
                pipe_id=pipe_id,
                status_callback=status_callback
            )
            logger.debug(f"Final Goal Result :\n{final_result.model_dump()}")
            if self.return_engine_result:
                final_result.engine_result = results
            return final_result
        else:
            engine_result = None
            if self.return_engine_result:
                engine_result = results
            return GoalResult(
                name=self.name,
                agent_id=self.agent_id,
                content=results,
                verify_goal=False,
                is_goal_satisfied=None,
                engine_result=engine_result
            )

    async def execute(
            self,
            *,
            query_instruction: str,
            pipe_id: str | None = None,
            pre_result: str | None = None,
            old_memory: list[dict] | None = None,
            verify_goal: bool = True,
            stop_if_goal_not_satisfied: bool = False,
            conversation_id: str | None = None,
            status_callback: StatusCallback | None = None
    ) -> GoalResult | None:
        """
        Executes the specified query instruction to achieve a defined goal.

        This method processes the `query_instruction`, potentially utilizing any
        pre-existing results provided through the `pre_result` parameter. The execution
        is designed to perform the necessary operations to fulfill the goal associated
        with the instruction and return a structured result indicating the outcome.

        Args:
            query_instruction: A string representing the instruction or query that defines the goal to be achieved.
                This should be a clear and actionable statement that the method can execute.
            pre_result: An optional pre-computed result or state to be used during the execution.
                Defaults to `None` if not provided.
            old_memory: An optional previous context of the user's instruction
            verify_goal: Option to enable or disable goal verification after agent execution. Default `True`
            stop_if_goal_not_satisfied: A flag indicating whether to stop processing if the goal is not satisfied.
                When set to True, the engine operation will halt if the defined goal is not met,
                preventing any further actions. Defaults to False, allowing the process to continue regardless
                of goal satisfaction.
            conversation_id: A string representing the unique identifier of the conversation.
            status_callback: This optional status call back method helps enhance user experience to get live updates of
                agents executions

        Returns:
            GoalResult | None
                An instance of the GoalResult class indicating the outcome of the execution. This result may include
                details about the success or failure of the operation, along with relevant data. If the
                execution cannot be completed or if an error occurs, the method may return `None`.
        """

        _goal_result = None

        if not self.llm:
            verify_goal = False

        # Callback: agent execution started
        if status_callback:
            await _maybe_await(status_callback(
                event="agent_execute_start",
                pipe_id=pipe_id,
                agent_id=self.agent_id,
                agent=self.name,
                query=query_instruction,
                conversation_id=conversation_id
            ))

        for _retry in range(1, self.max_retry + 1):
            # Callback: retry started
            if status_callback:
                await _maybe_await(status_callback(
                    event=f"agent_execute",
                    pipe_id=pipe_id,
                    agent_id=self.agent_id,
                    agent=self.name,
                    retry=_retry,
                    conversation_id=conversation_id,
                ))

            logger.info(f"Agent `{self.name}` retry {_retry}")
            try:
                _goal_result = await self._execute(
                    query_instruction=query_instruction,
                    pre_result=pre_result,
                    old_memory=old_memory,
                    pipe_id=pipe_id,
                    verify_goal=verify_goal,
                    conversation_id=conversation_id,
                    status_callback=status_callback
                )

                # Callback: iteration complete
                if status_callback:
                    await _maybe_await(status_callback(
                        event="agent_iteration_complete",
                        pipe_id=pipe_id,
                        agent_id=self.agent_id,
                        agent=self.name,
                        retry=_retry,
                        goal_result=_goal_result,
                        conversation_id=conversation_id
                    ))

                if not verify_goal or _goal_result.is_goal_satisfied:
                    if status_callback:
                        await _maybe_await(status_callback(
                            event="agent_goal_satisfied",
                            pipe_id=pipe_id,
                            agent_id=self.agent_id,
                            agent=self.name,
                            retry=_retry,
                            goal_result=_goal_result,
                            conversation_id=conversation_id
                        ))
                    return _goal_result

                elif _goal_result.is_goal_satisfied is False and stop_if_goal_not_satisfied:
                    if status_callback:
                        await _maybe_await(status_callback(
                            event="agent_goal_unsatisfied_stopped",
                            pipe_id=pipe_id,
                            agent_id=self.agent_id,
                            agent=self.name,
                            retry=_retry,
                            goal_result=_goal_result,
                            conversation_id=conversation_id
                        ))

                    raise StopSuperAgentX(
                        message='SuperAgentX stopped forcefully since `stop` flag has been set!',
                        goal_result=_goal_result
                    )
            except StopSuperAgentX:
                raise
            except Exception as e:
                logger.exception(f"Error executing agent `{self.name}` on retry {_retry}: {e}")
                if status_callback:
                    await _maybe_await(status_callback(
                        event="agent_execute_error",
                        pipe_id=pipe_id,
                        agent_id=self.agent_id,
                        agent=self.name,
                        retry=_retry,
                        error=str(e),
                        conversation_id=conversation_id
                    ))

        logger.warning(f"Done engine `{self.name}` max retry {self.max_retry}!")

        # Callback: execution finished (max retries reached)
        if status_callback:
            await _maybe_await(status_callback(
                event="agent_execute_complete",
                pipe_id=pipe_id,
                agent_id=self.agent_id,
                agent=self.name,
                goal_result=_goal_result,
                conversation_id=conversation_id
            ))

        return _goal_result
