"""
Iterative Planner Agent - works towards an objective using sub-agents
"""

import asyncio
from typing import Any, Dict, List, Optional, Tuple, Type

from mcp import Tool
from mcp.types import TextContent

from fast_agent.agents.agent_types import AgentConfig, AgentType
from fast_agent.agents.llm_agent import LlmAgent
from fast_agent.agents.workflow.orchestrator_models import (
    Plan,
    PlanningStep,
    PlanResult,
    Step,
    TaskWithResult,
    format_plan_result,
    format_step_result_text,
)
from fast_agent.core.exceptions import AgentConfigError
from fast_agent.core.logging.logger import get_logger
from fast_agent.core.prompt import Prompt
from fast_agent.interfaces import AgentProtocol, ModelT
from fast_agent.types import PromptMessageExtended, RequestParams
from fast_agent.workflow_telemetry import (
    NoOpPlanTelemetryProvider,
    PlanEntry,
    PlanEntryStatus,
    PlanTelemetryProvider,
)

logger = get_logger(__name__)


ITERATIVE_PLAN_SYSTEM_PROMPT_TEMPLATE = """
You are an expert planner, able to Orchestrate complex tasks by breaking them down in to
manageable steps, and delegating tasks to Agents.

You work iteratively - given an Objective, you consider the current state of the plan,
decide the next step towards the goal. You document those steps and create clear instructions
for execution by the Agents, being specific about what you need to know to assess task completion. 

NOTE: A 'Planning Step' has a description, and a list of tasks that can be delegated 
and executed in parallel.

Agents have a 'description' describing their primary function, and a set of 'skills' that
represent Tools they can use in completing their function.

The following Agents are available to you:

{{agents}}

You must specify the Agent name precisely when generating a Planning Step. 

"""


ITERATIVE_PLAN_PROMPT_TEMPLATE2 = """

```
<fastagent:data>
    <fastagent:progress>
{plan_result}
    </fastagent:progress>

    <fastagent:status>
{plan_status}
{iterations_info}
    </fastagent:status>
</fastagent:data>
```

The overall objective is:

```
<fastagent:objective>
    {objective}
</fastagent:objective>
```

Produce the next step in the plan to complete the Objective. 

Consider the previous steps and results, and decide what needs to be done next.

Set "is_complete" to true when ANY of these conditions are met:
1. The objective has been achieved in full or substantively
2. The remaining work is minor or trivial compared to what's been accomplished
3. The plan has gathered sufficient information to answer the original request
4. The plan has no feasible way of completing the objective.

Only set is_complete to `true` if there are no outstanding tasks. 

Be decisive - avoid excessive planning steps that add little value. It's better to complete a plan early than to continue with marginal improvements. 

Focus on the meeting the core intent of the objective.

"""


DELEGATED_TASK_TEMPLATE = """
You are a component in an orchestrated workflow to achieve an objective.

The overall objective is:

```
<fastagent:objective>
    {objective}
</fastagent:objective>
```

Previous context in achieving the objective is below:

```
<fastagent:context>
    {context}
</fastagent:context>
```

Your job is to accomplish the "task" specified in `<fastagent:task>`. The overall objective and
previous context is supplied to inform your approach.

```
<fastagent:task>
   {task}
</fastagent:task>
```

Provide a direct, concise response on completion that makes it simple to assess the status of
the overall plan.
"""


PLAN_RESULT_TEMPLATE = """
The below shows the results of running a plan to meet the specified objective.

```
<fastagent:plan-results>
{plan_result}
</fastagent:plan-results>
```

The plan was stopped because {termination_reason}.

Provide a summary of the tasks completed and their outcomes to complete the Objective.
Use markdown formatting.

If the plan was marked as incomplete but the maximum number of iterations was reached,
make sure to state clearly what was accomplished and what remains to be done.


Complete the plan by providing an appropriate answer for the original objective. Provide a Mermaid diagram
(in code fences) showing the plan steps and their relationships, if applicable. Do not use parentheses in labels.
"""


class IterativePlanner(LlmAgent):
    """
    An agent that implements the orchestrator workflow pattern.

    Dynamically creates execution plans and delegates tasks
    to specialized worker agents, synthesizing their results into a cohesive output.
    Supports both full planning and iterative planning modes.
    """

    @property
    def agent_type(self) -> AgentType:
        """Return the type of this agent."""
        return AgentType.ITERATIVE_PLANNER

    def __init__(
        self,
        config: AgentConfig,
        agents: List[AgentProtocol],
        plan_iterations: int = -1,
        context: Optional[Any] = None,
        **kwargs,
    ) -> None:
        """
        Initialize an OrchestratorAgent.

        Args:
            config: Agent configuration or name
            agents: List of specialized worker agents available for task execution
            plan_type: Planning mode ("full" or "iterative")
            context: Optional context object
            **kwargs: Additional keyword arguments to pass to BaseAgent
        """
        if not agents:
            raise AgentConfigError("At least one worker agent must be provided")

        # Store agents by name for easier lookup
        self.agents: Dict[str, AgentProtocol] = {}
        for agent in agents:
            agent_name = agent.name
            self.agents[agent_name] = agent

        # Extract plan_type from kwargs before passing to parent
        kwargs.pop("plan_type", "full")
        super().__init__(config, context=context, **kwargs)

        self.plan_iterations = plan_iterations
        self._plan_telemetry_provider: PlanTelemetryProvider = NoOpPlanTelemetryProvider()

    @property
    def plan_telemetry(self) -> PlanTelemetryProvider:
        """Telemetry provider for emitting plan updates."""
        return self._plan_telemetry_provider

    @plan_telemetry.setter
    def plan_telemetry(self, provider: PlanTelemetryProvider | None) -> None:
        if provider is None:
            provider = NoOpPlanTelemetryProvider()
        self._plan_telemetry_provider = provider

    async def initialize(self) -> None:
        """Initialize the orchestrator agent and worker agents."""
        # Initialize all worker agents first if not already initialized
        for agent_name, agent in self.agents.items():
            if not getattr(agent, "initialized", False):
                logger.debug(f"Initializing agent: {agent_name}")
                await agent.initialize()

        # Format agent information using agent cards with XML formatting
        agent_descriptions = []
        for agent_name, agent in self.agents.items():
            agent_card = await agent.agent_card()
            # Format as XML for better readability in prompts
            xml_formatted = self._format_agent_card_as_xml(agent_card)
            agent_descriptions.append(xml_formatted)

        agents_str = "\n".join(agent_descriptions)

        # Replace {{agents}} placeholder in the system prompt template
        system_prompt = self.config.instruction.replace("{{agents}}", agents_str)

        # Update the config instruction with the formatted system prompt
        self.instruction = system_prompt
        # Initialize the base agent with the updated system prompt
        await super().initialize()

        self.initialized = True

    async def shutdown(self) -> None:
        """Shutdown the orchestrator agent and worker agents."""
        await super().shutdown()

        # Shutdown all worker agents
        for agent_name, agent in self.agents.items():
            try:
                await agent.shutdown()
            except Exception as e:
                logger.warning(f"Error shutting down agent {agent_name}: {str(e)}")

    async def generate_impl(
        self,
        messages: List[PromptMessageExtended],
        request_params: RequestParams | None = None,
        tools: List[Tool] | None = None,
    ) -> PromptMessageExtended:
        """
        Execute an orchestrated plan to process the input.

        Args:
            normalized_messages: Already normalized list of PromptMessageExtended
            request_params: Optional request parameters

        Returns:
            The final synthesized response from the orchestration
        """
        # Extract user request
        objective = messages[-1].all_text() if messages else ""
        plan_result = await self._execute_plan(objective, request_params)
        # Return the result
        return PromptMessageExtended(
            role="assistant",
            content=[TextContent(type="text", text=plan_result.result or "No result available")],
        )

    async def structured_impl(
        self,
        messages: List[PromptMessageExtended],
        model: Type[ModelT],
        request_params: Optional[RequestParams] = None,
    ) -> Tuple[ModelT | None, PromptMessageExtended]:
        """
        Execute an orchestration plan and parse the result into a structured format.

        Args:
            messages: List of messages to process
            model: Pydantic model to parse the response into
            request_params: Optional request parameters

        Returns:
            The parsed final response, or None if parsing fails
        """
        # Generate orchestration result
        response = await self.generate_impl(messages, request_params)

        # Try to parse the response into the specified model
        try:
            result_text = response.last_text() or "<no text>"
            prompt_message = PromptMessageExtended(
                role="user", content=[TextContent(type="text", text=result_text)]
            )
            assert self._llm
            return await self._llm.structured([prompt_message], model, request_params)
        except Exception as e:
            logger.warning(f"Failed to parse orchestration result: {str(e)}")
            return None, Prompt.assistant(f"Failed to parse orchestration result: {str(e)}")

    async def _execute_plan(
        self, objective: str, request_params: RequestParams | None
    ) -> PlanResult:
        """
        Execute a plan to achieve the given objective.

        Args:
            objective: The objective to achieve
            request_params: Request parameters for execution

        Returns:
            PlanResult containing execution results and final output
        """

        objective_met: bool = False
        terminate_plan: str | None = None
        plan_result = PlanResult(objective=objective, step_results=[])

        # Track plan entries for telemetry
        plan_entries: list[PlanEntry] = []

        while not objective_met and not terminate_plan:
            next_step: PlanningStep | None = await self._get_next_step(
                objective, plan_result, request_params
            )

            if None is next_step:
                terminate_plan = "Failed to generate plan, terminating early"
                logger.error("Failed to generate next step, terminating plan early")
                break

            assert next_step  # lets keep the indenting manageable!

            if next_step.is_complete:
                objective_met = True
                terminate_plan = "Plan completed successfully"
                # Mark all entries as completed
                for entry in plan_entries:
                    entry.status = "completed"
                await self.plan_telemetry.update_plan(plan_entries)
                break

            plan = Plan(steps=[next_step], is_complete=next_step.is_complete)
            invalid_agents = self._validate_agent_names(plan)
            if invalid_agents:
                logger.error(f"Plan contains invalid agent names: {', '.join(invalid_agents)}")
                terminate_plan = (
                    f"Invalid agent names found ({', '.join(invalid_agents)}), terminating plan"
                )
                break

            # Add subtasks as plan entries for telemetry so clients see actionable work
            step_entries = self._build_plan_entries_for_step(next_step)
            plan_entries.extend(step_entries)
            await self.plan_telemetry.update_plan(plan_entries)

            # Mark subtasks as in progress while they execute
            self._set_plan_entries_status(step_entries, "in_progress")
            await self.plan_telemetry.update_plan(plan_entries)

            for step in plan.steps:  # this will only be one for iterative (change later)
                step_result = await self._execute_step(step, plan_result)
                plan_result.add_step_result(step_result)

                # Mark subtasks as completed
                self._set_plan_entries_status(step_entries, "completed")
                await self.plan_telemetry.update_plan(plan_entries)

            # Store plan in result
            plan_result.plan = plan

            if self.plan_iterations > 0:
                if len(plan_result.step_results) >= self.plan_iterations:
                    terminate_plan = f"Reached maximum number of iterations ({self.plan_iterations}), terminating plan"

        if not terminate_plan:
            terminate_plan = "Unknown termination reason"
        result_prompt = PLAN_RESULT_TEMPLATE.format(
            plan_result=format_plan_result(plan_result), termination_reason=terminate_plan
        )

        # Generate final synthesis
        final_message = await self._planner_generate_str(result_prompt, request_params)
        plan_result.result = final_message.last_text() or "No final message generated"
        await self.show_assistant_message(final_message)
        return plan_result

    async def _execute_step(self, step: Step, previous_result: PlanResult) -> Any:
        """
        Execute a single step from the plan.

        Args:
            step: The step to execute
            previous_result: Results of the plan execution so far
            request_params: Request parameters

        Returns:
            Result of executing the step
        """
        from fast_agent.agents.workflow.orchestrator_models import StepResult

        # Initialize step result
        step_result = StepResult(step=step, task_results=[])

        # Format context for tasks
        context = format_plan_result(previous_result)

        # Group tasks by agent and execute different agents in parallel
        from collections import defaultdict

        tasks_by_agent = defaultdict(list)
        for task in step.tasks:
            tasks_by_agent[task.agent].append(task)

        async def execute_agent_tasks(agent_name: str, agent_tasks: List) -> List[TaskWithResult]:
            """Execute all tasks for a single agent sequentially (preserves history)"""
            agent = self.agents.get(agent_name)
            assert agent is not None

            results = []
            for task in agent_tasks:
                try:
                    task_description = DELEGATED_TASK_TEMPLATE.format(
                        objective=previous_result.objective, task=task.description, context=context
                    )
                    result = await agent.generate(
                        [
                            PromptMessageExtended(
                                role="user",
                                content=[TextContent(type="text", text=task_description)],
                            )
                        ]
                    )

                    task_model = task.model_dump()
                    results.append(
                        TaskWithResult(
                            description=task_model["description"],
                            agent=task_model["agent"],
                            result=result.last_text() or "<missing response>",
                        )
                    )
                except Exception as e:
                    logger.error(f"Error executing task: {str(e)}")
                    task_model = task.model_dump()
                    results.append(
                        TaskWithResult(
                            description=task_model["description"],
                            agent=task_model["agent"],
                            result=f"ERROR: {str(e)}",
                        )
                    )
            return results

        # Execute different agents in parallel, tasks within each agent sequentially
        agent_futures = [
            execute_agent_tasks(agent_name, agent_tasks)
            for agent_name, agent_tasks in tasks_by_agent.items()
        ]

        all_results = await asyncio.gather(*agent_futures)
        task_results = [result for agent_results in all_results for result in agent_results]

        # Add all task results to step result
        for task_result in task_results:
            step_result.add_task_result(task_result)

        # Format step result
        step_result.result = format_step_result_text(step_result)
        return step_result

    async def _get_next_step(
        self, objective: str, plan_result: PlanResult, request_params: RequestParams | None
    ) -> PlanningStep | None:
        """
        Generate just the next step for iterative planning.

        Args:
            objective: The objective to achieve
            plan_result: Current plan execution state
            request_params: Request parameters

        Returns:
            Next step to execute, or None if parsing fails
        """

        # Determine plan status
        if plan_result.is_complete:
            plan_status = "Plan Status: Complete"
        elif plan_result.step_results:
            plan_status = "Plan Status: In Progress"
        else:
            plan_status = "Plan Status: Not Started"

        # Calculate iteration information

        if self.plan_iterations > 0:
            max_iterations = self.plan_iterations
            current_iteration = len(plan_result.step_results)
            iterations_remaining = max_iterations - current_iteration
            iterations_info = (
                f"Planning Budget: {iterations_remaining} of {max_iterations} iterations remaining"
            )
        else:
            iterations_info = "Iterating until objective is met."

        # Format the planning prompt
        prompt = ITERATIVE_PLAN_PROMPT_TEMPLATE2.format(
            objective=objective,
            plan_result=format_plan_result(plan_result),
            plan_status=plan_status,
            iterations_info=iterations_info,
        )

        request_params = self._merge_request_params(
            request_params, RequestParams(systemPrompt=self.instruction)
        )

        # Get structured response from LLM
        try:
            plan_msg = PromptMessageExtended(
                role="user", content=[TextContent(type="text", text=prompt)]
            )
            assert self._llm
            self.show_user_message(plan_msg)
            next_step, raw_response = await self._llm.structured(
                [plan_msg], PlanningStep, request_params
            )
            await self.show_assistant_message(raw_response)
            return next_step
        except Exception as e:
            logger.error(f"Failed to parse next step: {str(e)}")
            return None

    def _validate_agent_names(self, plan: Plan) -> List[str]:
        """
        Validate all agent names in a plan before execution.

        Args:
            plan: The plan to validate
        """
        invalid_agents = []

        for step in plan.steps:
            for task in step.tasks:
                if task.agent not in self.agents:
                    invalid_agents.append(task.agent)

        return invalid_agents

    @staticmethod
    def _format_agent_card_as_xml(agent_card) -> str:
        """
        Format an agent card as XML for display in prompts.

        This creates a structured XML representation that's more readable than JSON
        and includes all relevant agent information in a hierarchical format.

        Args:
            agent_card: The AgentCard object

        Returns:
            XML formatted agent information
        """
        xml_parts = [f'<fastagent:agent name="{agent_card.name}">']

        # Add description if available
        if agent_card.description:
            xml_parts.append(
                f"  <fastagent:description>{agent_card.description}</fastagent:description>"
            )

        # Add skills if available
        if hasattr(agent_card, "skills") and agent_card.skills:
            xml_parts.append("  <fastagent:skills>")
            for skill in agent_card.skills:
                xml_parts.append(f'    <fastagent:skill name="{skill.name}">')
                if hasattr(skill, "description") and skill.description:
                    xml_parts.append(
                        f"      <fastagent:description>{skill.description}</fastagent:description>"
                    )
                xml_parts.append("    </fastagent:skill>")
            xml_parts.append("  </fastagent:skills>")

        xml_parts.append("</fastagent:agent>")

        return "\n".join(xml_parts)

    async def _planner_generate_str(
        self, message: str, request_params: RequestParams | None
    ) -> PromptMessageExtended:
        """
        Generate string response from the orchestrator's own LLM.

        Args:
            message: Message to send to the LLM
            request_params: Request parameters

        Returns:
            String response from the LLM
        """
        # Create prompt message
        prompt = PromptMessageExtended(
            role="user", content=[TextContent(type="text", text=message)]
        )
        assert self._llm, "LLM must be initialized before generating text"
        return await self._llm.generate([prompt], request_params)

    @staticmethod
    def _build_plan_entries_for_step(step: PlanningStep | Step) -> list[PlanEntry]:
        """
        Build plan entries for telemetry using individual step tasks when available.
        """
        if step.tasks:
            entries = []
            for task in step.tasks:
                agent_name = getattr(task, "agent", None) or "Unknown agent"
                entries.append(
                    PlanEntry(
                        content=f"{agent_name}: {task.description}",
                        priority="high",
                        status="pending",
                    )
                )
            return entries

        # Fallback to a single entry if the planner supplied no explicit tasks
        return [
            PlanEntry(
                content=step.description,
                priority="high",
                status="pending",
            )
        ]

    @staticmethod
    def _set_plan_entries_status(entries: list[PlanEntry], status: PlanEntryStatus) -> None:
        """Helper to bulk update plan entry statuses."""
        for entry in entries:
            entry.status = status
