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.channels.base import HumanApprovalChannel
from superagentx.channels.console_channel import ConsoleApprovalChannel
from superagentx.db_store import StorageAdapter
from superagentx.handler.base import BaseHandler
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
from superagentx.utils.observability.span_decorator import agent_span

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,
            tool: BaseHandler | None = None,
            tool_args: dict[str, Any] | None = None,
            output_format: str | None = None,
            max_retry: int = 5,
            human_approval: bool = False,
            approval_channel: HumanApprovalChannel = None,
            return_engine_result: bool = False
    ):
        """
        Initializes a new Agent instance.

        This constructor configures an agent capable of interacting with a large
        language model (LLM), executing tasks through one or more engines, and
        optionally invoking tools or requiring human approval during execution.
        The agent's behavior can be customized via goals, roles, prompt templates,
        and execution controls.

        Args:
            goal: The primary objective the agent is intended to accomplish.
                Used to guide reasoning and task execution.
            role: The role or persona the agent should assume when operating.
                Typically influences prompt construction and behavior.
            llm: The LLM client used to communicate with the underlying large
                language model.
            prompt_template: Template defining the structure and format of prompts
                sent to the LLM.
            agent_id: A unique identifier for the agent. If not provided, a UUID
                may be generated automatically.
            name: A human-readable name for the agent, useful for logging,
                debugging, or display purposes.
            description: Optional descriptive text providing additional context
                about the agent’s purpose or capabilities.
            engines: A list of engines (or nested lists of engines) available to
                the agent for task execution. Engines may support different
                capabilities such as browsing, task orchestration, or tool usage.
            tool: An optional tool handler that can be invoked by the agent when
                no engines are explicitly provided.
            tool_args: Optional arguments passed to the tool handler when invoked.
            output_format: Specifies the desired format of the agent’s output
                (e.g., text, JSON, or a custom schema).
            max_retry: Maximum number of retry attempts for recoverable failures
                during execution. Defaults to 5.
            human_approval: Whether certain actions require explicit human
                approval before proceeding.
            approval_channel: The channel or mechanism used to request and receive
                human approval when enabled.
            return_engine_result: Whether to return raw engine execution results
                instead of (or in addition to) post-processed agent output.
        """
        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.tool = tool
        self.tool_args = tool_args
        self.output_format = output_format
        self.max_retry = max_retry if max_retry >= 1 else 1
        self.return_engine_result = return_engine_result
        self.human_approval = human_approval
        self.approval_channel = approval_channel or ConsoleApprovalChannel()
        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
                        )
            return None
        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,
            storage: StorageAdapter = 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,
            "storage": storage,
            "status_callback": status_callback
        }
        if not self.engines and self.tool: # If pass tool via Agent
            logger.debug(f'Engine(s) is empty')
            engine = Engine(handler=self.tool, llm=self.llm, prompt_template=self.prompt_template)
            await self.add(engine)

        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 _request_human_approval(
            self,
            *,
            query: str,
            pre_result,
            pipe_id=None,
            conversation_id=None
    ) -> bool:
        return await self.approval_channel.request_approval(
            agent_id=self.agent_id,
            agent_name=self.name,
            query=query,
            pre_result=pre_result,
            pipe_id=pipe_id,
            conversation_id=conversation_id
        )

    @agent_span
    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,
            storage: StorageAdapter = 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.
            pipe_id: Pipe interface execution id.
            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.
            storage: An optional storage adapter to store Agentic Workflow states. This is useful, if a human
                    approvals are enabled to wait and act after a human action is being performed.
            status_callback: This optional status call back method helps enhance user experience to get live updates of
                agents executions

            Executes the agent with optional human approval.

                Status rules:
                    - No human approval → COMPLETED / ERROR
                    - With human approval → APPROVED / REJECTED

        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

        # ------------------------------------------------------------------
        # HUMAN APPROVAL (FINAL AGENT STATE)
        # ------------------------------------------------------------------
        if self.human_approval:
            if storage:
                await storage.update_pipe_status(pipe_id, "Waiting-for-Approval")

            approved = await self._request_human_approval(
                query=query_instruction,
                pre_result=pre_result,
                pipe_id=pipe_id,
                conversation_id=conversation_id
            )

            if not approved:
                if storage:
                    await storage.mark_agent_completed(
                        pipe_id=pipe_id,
                        agent_id=self.agent_id,
                        agent_name=self.name,
                        input_content=query_instruction,
                        status="REJECTED"
                    )
                    await storage.update_pipe_status(
                        pipe_id,
                        "Failed",
                        error="Human rejected execution"
                    )

                if status_callback:
                    await _maybe_await(status_callback(
                        event="human_status_rejected",
                        pipe_id=pipe_id,
                        agent_id=self.agent_id,
                        agent=self.name,
                        query=query_instruction,
                        conversation_id=conversation_id
                    ))

                raise StopSuperAgentX(
                    message="Execution rejected by human approval.",
                    goal_result=GoalResult(
                        name=self.name,
                        agent_id=self.agent_id,
                        content="Rejected by human",
                        approved_status="REJECTED",
                        is_goal_satisfied=False
                    )
                )

            # Approved is FINAL agent status
            if storage:
                await storage.mark_agent_completed(
                    pipe_id=pipe_id,
                    agent_id=self.agent_id,
                    agent_name=self.name,
                    input_content=query_instruction,
                    status="APPROVED"
                )
                await storage.update_pipe_status(pipe_id, "In-Progress")

            if status_callback:
                await _maybe_await(status_callback(
                    event="human_status_approved",
                    pipe_id=pipe_id,
                    agent_id=self.agent_id,
                    agent=self.name,
                    query=query_instruction,
                    conversation_id=conversation_id
                ))

        # ------------------------------------------------------------------
        # EXECUTION START CALLBACK
        # ------------------------------------------------------------------
        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
            ))

        # ------------------------------------------------------------------
        # EXECUTION LOOP
        # ------------------------------------------------------------------
        for retry in range(1, self.max_retry + 1):
            try:
                if status_callback:
                    await _maybe_await(status_callback(
                        event="agent_execute",
                        pipe_id=pipe_id,
                        agent_id=self.agent_id,
                        agent=self.name,
                        retry=retry,
                        conversation_id=conversation_id
                    ))

                _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,
                    storage=storage,
                    conversation_id=conversation_id,
                    status_callback=status_callback
                )

                # Store COMPLETED only when NO human approval
                if storage and not self.human_approval:
                    await storage.mark_agent_completed(
                        pipe_id=pipe_id,
                        agent_id=self.agent_id,
                        agent_name=self.name,
                        input_content=query_instruction,
                        goal_result=_goal_result,
                        status="COMPLETED"
                    )

                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,
                            goal_result=_goal_result,
                            conversation_id=conversation_id
                        ))
                    return _goal_result

                if stop_if_goal_not_satisfied:
                    raise StopSuperAgentX(
                        message="Goal not satisfied. Stopping execution.",
                        goal_result=_goal_result
                    )

            except StopSuperAgentX:
                raise

            except Exception as e:
                logger.exception(f"Agent `{self.name}` failed on retry {retry}: {e}")

                # Store ERROR only when NO human approval
                if storage and not self.human_approval:
                    await storage.mark_agent_completed(
                        pipe_id=pipe_id,
                        agent_id=self.agent_id,
                        agent_name=self.name,
                        input_content=query_instruction,
                        goal_result=_goal_result,
                        status="ERROR"
                    )

                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
                    ))

        # ------------------------------------------------------------------
        # EXECUTION FINISHED (MAX RETRIES)
        # ------------------------------------------------------------------
        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
