"""Plan-based orchestration pattern."""

from __future__ import annotations

from typing import Any, Optional, Sequence

from pydantic import BaseModel, ConfigDict, Field

from agentbyte.agents.base import BaseAgent
from agentbyte.agents.types import AgentResponse, Usage
from agentbyte.llm.base import BaseChatCompletionClient
from agentbyte.messages import Message, UserMessage
from agentbyte.termination import BaseTermination
from agentbyte.types import StopMessage

from .base import BaseOrchestrator
from .policies import (
    AgentHistoryPolicy,
    FinalResultPolicy,
    HistoryContextPolicy,
    ReviewGatePolicy,
)


class StepProgressEvaluation(BaseModel):
    """Structured evaluation of whether a plan step succeeded."""

    model_config = ConfigDict(extra="forbid")

    step_completed: bool = Field(
        ..., description="Whether the step was successfully completed"
    )
    failure_reason: str = Field(
        ..., description="Why the step failed, or 'None' if it succeeded"
    )
    confidence_score: float = Field(
        ...,
        ge=0.0,
        le=1.0,
        description="Confidence in the evaluation from 0.0 to 1.0",
    )
    suggested_improvements: list[str] = Field(
        ...,
        description="Specific suggestions to improve a retry attempt",
    )


class PlanStep(BaseModel):
    """Single step in an execution plan."""

    model_config = ConfigDict(extra="forbid")

    task: str = Field(..., description="Clear, actionable task description")
    agent_name: str = Field(..., description="Name of the agent for this step")
    reasoning: str = Field(..., description="Why this agent was chosen")


class ExecutionPlan(BaseModel):
    """Ordered execution plan produced by the planner model."""

    model_config = ConfigDict(extra="forbid")

    steps: list[PlanStep] = Field(..., description="Ordered list of execution steps")


class PlanBasedOrchestrator(BaseOrchestrator):
    """
    Orchestrator that plans work up front, evaluates each step, and retries when useful.
    """

    def __init__(
        self,
        agents: Sequence[BaseAgent],
        termination: BaseTermination,
        model_client: BaseChatCompletionClient,
        max_iterations: int = 50,
        max_step_retries: int = 3,
        history_policy: HistoryContextPolicy | None = None,
        review_gate_policy: ReviewGatePolicy | None = None,
        final_result_policy: FinalResultPolicy | None = None,
    ):
        super().__init__(
            agents,
            termination,
            max_iterations,
            history_policy=history_policy,
            final_result_policy=final_result_policy,
        )
        self.model_client = model_client
        self.max_step_retries = max(0, max_step_retries)
        self.review_gate_policy = review_gate_policy or ReviewGatePolicy()
        self._validate_review_gate_policy()

        self.execution_plan: Optional[ExecutionPlan] = None
        self.current_step_index = 0
        self.current_step_retry_count = 0
        self.initial_task: Optional[str] = None

        self.step_attempts: dict[int, list[AgentResponse]] = {}
        self.step_results: dict[int, AgentResponse] = {}
        self.retry_instructions: dict[int, str] = {}
        self._failed_step_indices: set[int] = set()
        self._orchestrator_usage_stats: list[Usage] = []
        self._agent_capabilities_cache: Optional[str] = None
        self._review_revision_rounds = 0

    async def select_next_agent(self) -> BaseAgent:
        """Select the agent assigned to the current plan step."""
        if self.execution_plan is None:
            self.initial_task = self._extract_initial_task()
            self.execution_plan = await self.create_plan(self.initial_task)

        current_step = self._get_current_step()
        if current_step is None:
            return self.agents[0]

        return self._find_agent_by_name(current_step.agent_name)

    async def prepare_context_for_agent(
        self, agent: BaseAgent
    ) -> UserMessage | list[Message]:
        """Provide policy-shaped history plus the current plan step instructions."""
        current_step = self._get_current_step()
        base_context = self._build_agent_context_from_history(agent)
        if current_step is None:
            return base_context

        step_instruction = UserMessage(
            content=self._format_step_task(current_step),
            source="plan_orchestrator",
        )

        if isinstance(base_context, list):
            context = list(base_context)
            context.append(step_instruction)
            self._last_context_size = len(context)
            self._last_context_was_text = False
            return context

        context = UserMessage(
            content=f"{base_context.content}\n\n{step_instruction.content}",
            source=base_context.source,
        )
        self._last_context_size = 1
        self._last_context_was_text = True
        return context

    async def update_shared_state(self, result: AgentResponse) -> None:
        """Track step attempts and advance or retry based on evaluation policies."""
        if self._last_context_was_text:
            candidate_messages = result.messages[1:] if len(result.messages) > 1 else []
        else:
            candidate_messages = result.messages[self._last_context_size :]

        new_messages = [
            message
            for message in candidate_messages
            if not isinstance(message, UserMessage)
        ]
        self.shared_messages.extend(new_messages)

        current_step = self._get_current_step()
        if current_step is None:
            return

        attempts = self.step_attempts.setdefault(self.current_step_index, [])
        attempts.append(result)

        evaluation = await self.evaluate_step_progress(current_step, result)
        if evaluation.step_completed:
            self.step_results[self.current_step_index] = result
            self.retry_instructions.pop(self.current_step_index, None)
            if self._is_review_gate_step(current_step):
                self._review_revision_rounds = 0
            self.current_step_index += 1
            self.current_step_retry_count = 0
            return

        if self._should_insert_revision_loop(current_step):
            self._review_revision_rounds += 1
            reviewer_feedback = self._extract_agent_output(result)
            insertion_index = self.current_step_index + 1
            assert self.execution_plan is not None
            self.execution_plan.steps[insertion_index:insertion_index] = [
                PlanStep(
                    task=(
                        "Revise the draft based on the reviewer feedback below and "
                        "produce an updated final draft:\n\n"
                        f"{reviewer_feedback}"
                    ),
                    agent_name=self.review_gate_policy.revision_agent_name,
                    reasoning=(
                        "The reviewer requested revisions, so the revision agent "
                        "should update the draft."
                    ),
                ),
                PlanStep(
                    task=(
                        "Review the revised draft. Respond with "
                        f"{self.review_gate_policy.approval_signal} only if it is ready. "
                        "Otherwise provide concise feedback on what still needs improvement."
                    ),
                    agent_name=self.review_gate_policy.reviewer_agent_name,
                    reasoning="The reviewer is the final approval gate after the revision pass.",
                ),
            ]
            self.current_step_index += 1
            self.current_step_retry_count = self._review_revision_rounds
            return

        self.current_step_retry_count += 1
        if self.current_step_retry_count <= self.max_step_retries:
            self.retry_instructions[self.current_step_index] = (
                self._create_retry_instructions(current_step, evaluation)
            )
            return

        self._failed_step_indices.add(self.current_step_index)
        self.retry_instructions.pop(self.current_step_index, None)
        self.current_step_index += 1
        self.current_step_retry_count = 0

    async def create_plan(self, task: str) -> ExecutionPlan:
        """Create a structured execution plan with model-backed fallback behavior."""
        planning_prompt = self._build_planning_prompt(task)

        try:
            result = await self.model_client.create(
                messages=[UserMessage(content=planning_prompt, source="planner")],
                output_format=ExecutionPlan,
            )
            self._orchestrator_usage_stats.append(result.usage)

            structured_output = result.structured_output
            if isinstance(structured_output, ExecutionPlan):
                return structured_output
            if isinstance(structured_output, dict):
                return ExecutionPlan.model_validate(structured_output)
        except Exception:
            return self._create_fallback_plan(task)

        return self._create_fallback_plan(task)

    async def evaluate_step_progress(
        self, step: PlanStep, result: AgentResponse
    ) -> StepProgressEvaluation:
        """Assess whether the current step was completed successfully."""
        agent_output = self._extract_agent_output(result)
        if self._is_review_gate_step(step):
            return self._evaluate_review_gate_step(agent_output)

        if not agent_output:
            return StepProgressEvaluation(
                step_completed=False,
                failure_reason="No meaningful output detected",
                confidence_score=0.9,
                suggested_improvements=[
                    "Provide clearer instructions",
                    "Include an example of the expected result",
                ],
            )

        evaluation_prompt = self._build_evaluation_prompt(step, agent_output)
        try:
            result_eval = await self.model_client.create(
                messages=[
                    UserMessage(content=evaluation_prompt, source="step_evaluator")
                ],
                output_format=StepProgressEvaluation,
            )
            self._orchestrator_usage_stats.append(result_eval.usage)

            structured_output = result_eval.structured_output
            if isinstance(structured_output, StepProgressEvaluation):
                return structured_output
            if isinstance(structured_output, dict):
                return StepProgressEvaluation.model_validate(structured_output)
        except Exception:
            return self._fallback_step_evaluation(agent_output)

        return self._fallback_step_evaluation(agent_output)

    def get_agent_capabilities_summary(self) -> str:
        """Cache the capabilities summary because planning reuses it."""
        if self._agent_capabilities_cache is None:
            self._agent_capabilities_cache = super().get_agent_capabilities_summary()
        return self._agent_capabilities_cache

    def _get_additional_usage_stats(self) -> list[Usage]:
        """Include planning and step-evaluation calls in final orchestration usage."""
        return list(self._orchestrator_usage_stats)

    def _get_pattern_stop_message(self) -> Optional[StopMessage]:
        """Stop cleanly once the full plan has been executed."""
        if self.execution_plan and self.current_step_index >= len(self.execution_plan.steps):
            return StopMessage(
                content="Execution plan completed",
                source="PlanBasedOrchestrator",
                metadata={
                    "steps_completed": len(self.step_results),
                    "steps_failed": len(self._failed_step_indices),
                },
            )
        return None

    def _get_pattern_metadata(self) -> dict[str, Any]:
        """Expose plan-aware runtime metadata."""
        base_metadata = super()._get_pattern_metadata()
        base_metadata.update(
            {
                "history_policy": self.history_policy.model_dump(),
                "review_gate_policy": self.review_gate_policy.model_dump(),
                "max_step_retries": self.max_step_retries,
                "steps_completed": len(self.step_results),
                "steps_failed": len(self._failed_step_indices),
                "total_retries": sum(
                    max(0, len(attempts) - 1) for attempts in self.step_attempts.values()
                ),
                "current_step_index": self.current_step_index,
                "current_step_retry_count": self.current_step_retry_count,
                "planner_model": self.model_client.model,
                "planner_calls": 1 if self.execution_plan else 0,
                "evaluation_calls": len(self.step_attempts),
            }
        )
        if self.execution_plan is not None:
            base_metadata["execution_plan"] = self.execution_plan.model_dump()
            base_metadata["total_steps"] = len(self.execution_plan.steps)
        return base_metadata

    def _reset_for_run(self) -> None:
        """Reset plan-specific runtime state."""
        super()._reset_for_run()
        self.execution_plan = None
        self.current_step_index = 0
        self.current_step_retry_count = 0
        self.initial_task = None
        self.step_attempts = {}
        self.step_results = {}
        self.retry_instructions = {}
        self._failed_step_indices = set()
        self._orchestrator_usage_stats = []
        self._agent_capabilities_cache = None
        self._review_revision_rounds = 0

    def _extract_initial_task(self) -> str:
        """Use the first message as the planning input when possible."""
        if not self.shared_messages:
            raise ValueError("No initial task found to create plan")
        return self.shared_messages[0].content

    def _get_current_step(self) -> Optional[PlanStep]:
        """Return the active step or None if the plan is complete."""
        if self.execution_plan is None:
            return None
        if self.current_step_index >= len(self.execution_plan.steps):
            return None
        return self.execution_plan.steps[self.current_step_index]

    def _build_planning_prompt(self, task: str) -> str:
        """Create the planner prompt."""
        return (
            "You are a helpful assistant that breaks down tasks into executable steps.\n\n"
            "Available agents and their capabilities:\n"
            f"{self.get_agent_capabilities_summary()}\n\n"
            f"User task: {task}\n\n"
            "Generate a concise execution plan. For each step:\n"
            "- Assign it to the agent best suited for that work\n"
            "- Provide a clear, actionable task description\n"
            "- Explain briefly why that agent was chosen\n\n"
            "Keep it focused. If only 2-3 steps are needed, that is fine."
        )

    def _build_evaluation_prompt(self, step: PlanStep, agent_output: str) -> str:
        """Create the prompt used to evaluate a step attempt."""
        return (
            "Evaluate whether the following plan step was completed successfully.\n\n"
            f"Step task: {step.task}\n"
            f"Assigned agent: {step.agent_name}\n"
            f"Assignment reasoning: {step.reasoning}\n\n"
            "Agent output:\n"
            f"{agent_output}\n\n"
            "Determine:\n"
            "1. Whether the step was completed successfully.\n"
            "2. The main failure reason if it was not.\n"
            "3. Confidence from 0.0 to 1.0.\n"
            "4. Specific suggestions for improvement if a retry would help.\n\n"
            "Treat meaningful progress toward the task as success even if the result is not perfect."
        )

    def _extract_agent_output(self, result: AgentResponse) -> str:
        """Collect non-user message content for evaluator prompts."""
        output_lines = [
            message.content.strip()
            for message in result.messages
            if not isinstance(message, UserMessage) and getattr(message, "content", "").strip()
        ]
        return "\n".join(output_lines).strip()

    def _fallback_step_evaluation(self, agent_output: str) -> StepProgressEvaluation:
        """Use deterministic heuristics when evaluator output is unavailable."""
        output_lower = agent_output.lower()
        has_meaningful_content = len(agent_output.strip()) > 20
        has_error_indicators = any(
            token in output_lower for token in ["error", "failed", "cannot", "unable", "sorry"]
        )
        has_positive_indicators = any(
            token in output_lower
            for token in ["completed", "found", "created", "analyzed", "wrote", "reviewed"]
        )

        if has_meaningful_content and (has_positive_indicators or not has_error_indicators):
            return StepProgressEvaluation(
                step_completed=True,
                failure_reason="None",
                confidence_score=0.7,
                suggested_improvements=[],
            )

        return StepProgressEvaluation(
            step_completed=False,
            failure_reason="Output suggests the task was not completed successfully",
            confidence_score=0.6,
            suggested_improvements=[
                "Provide clearer instructions",
                "Break the task into a smaller step",
                "Include an example of the expected output",
            ],
        )

    def _create_fallback_plan(self, task: str) -> ExecutionPlan:
        """Use a safe single-step plan if planning fails."""
        return ExecutionPlan(
            steps=[
                PlanStep(
                    task=f"Complete the task: {task}",
                    agent_name=self.agents[0].name,
                    reasoning="Single-step fallback plan",
                )
            ]
        )

    def _find_agent_by_name(self, agent_name: str) -> BaseAgent:
        """Find an agent by case-insensitive name, falling back safely."""
        target = agent_name.casefold().strip()
        for agent in self.agents:
            if agent.name.casefold() == target:
                return agent
        return self.agents[0]

    def _format_step_task(self, step: PlanStep) -> str:
        """Format the current step and any retry instructions."""
        task_lines = [
            f"PLAN STEP {self.current_step_index + 1}: {step.task}",
            f"Why you were chosen: {step.reasoning}",
        ]
        retry_instruction = self.retry_instructions.get(self.current_step_index)
        if retry_instruction:
            task_lines.extend(
                [
                    "",
                    f"RETRY ATTEMPT {self.current_step_retry_count + 1}:",
                    retry_instruction,
                ]
            )
        return "\n".join(task_lines)

    def _create_retry_instructions(
        self, step: PlanStep, evaluation: StepProgressEvaluation
    ) -> str:
        """Turn evaluator feedback into better retry instructions."""
        improvements = evaluation.suggested_improvements or [
            "Take a different approach while staying focused on the step goal."
        ]
        improvement_text = "\n".join(f"- {item}" for item in improvements)
        return (
            f"Previous attempt did not complete the step.\n"
            f"Failure reason: {evaluation.failure_reason}\n"
            f"Original step goal: {step.task}\n"
            "Improve the next attempt by addressing:\n"
            f"{improvement_text}"
        )

    def _evaluate_review_gate_step(self, agent_output: str) -> StepProgressEvaluation:
        """Evaluate review-gated steps via explicit approval signal matching."""
        approval_signal = self.review_gate_policy.approval_signal.casefold()
        normalized_lines = [
            line.strip().casefold()
            for line in agent_output.splitlines()
            if line.strip()
        ]
        if approval_signal in normalized_lines:
            return StepProgressEvaluation(
                step_completed=True,
                failure_reason="None",
                confidence_score=0.99,
                suggested_improvements=[],
            )

        return StepProgressEvaluation(
            step_completed=False,
            failure_reason=(
                "Reviewer did not explicitly approve the content. "
                f"The reviewer must respond with {self.review_gate_policy.approval_signal} "
                "when the final draft is ready."
            ),
            confidence_score=0.95,
            suggested_improvements=[
                "Revise the draft using the review feedback.",
                (
                    "Ask the reviewer to explicitly confirm readiness with "
                    f"{self.review_gate_policy.approval_signal} once resolved."
                ),
            ],
        )

    def _is_review_gate_step(self, step: PlanStep) -> bool:
        """Return whether the step belongs to the configured review gate."""
        if not self.review_gate_policy.enabled:
            return False
        return (
            step.agent_name.casefold().strip()
            == self.review_gate_policy.reviewer_agent_name.casefold().strip()
        )

    def _should_insert_revision_loop(self, step: PlanStep) -> bool:
        """Return whether a failed review step should trigger a revision loop."""
        return (
            self._is_review_gate_step(step)
            and self.execution_plan is not None
            and self._review_revision_rounds < self.review_gate_policy.max_revision_rounds
        )

    def _validate_review_gate_policy(self) -> None:
        """Ensure enabled review-gate policies reference existing agents."""
        if not self.review_gate_policy.enabled:
            return

        configured_names = {
            agent.name.casefold(): agent.name for agent in self.agents
        }
        reviewer_name = self.review_gate_policy.reviewer_agent_name.casefold()
        revision_name = self.review_gate_policy.revision_agent_name.casefold()

        if reviewer_name not in configured_names:
            raise ValueError(
                "ReviewGatePolicy.reviewer_agent_name must reference an agent in the pool"
            )
        if revision_name not in configured_names:
            raise ValueError(
                "ReviewGatePolicy.revision_agent_name must reference an agent in the pool"
            )

    def _default_history_policy(self) -> HistoryContextPolicy:
        """Use a recent structured message window by default."""
        return HistoryContextPolicy(
            default=AgentHistoryPolicy(
                window="recent",
                message_limit=5,
                format="messages",
            )
        )


__all__ = [
    "ExecutionPlan",
    "PlanBasedOrchestrator",
    "PlanStep",
    "StepProgressEvaluation",
]
