from typing import Any, Dict, List, Optional
from pygeai_orchestration.core.base import BasePattern, PatternConfig, PatternResult, PatternType
from pygeai_orchestration.core.common import Message, MessageRole, Conversation, State, StateStatus
import logging

logger = logging.getLogger("pygeai_orchestration")


class PlanStep:
    """
    Represents a single step in a planning execution.

    :param step_num: int - The step number in the overall plan.
    :param description: str - Description of what this step accomplishes.
    :param status: str - Current status: 'pending', 'completed', or 'failed'. Defaults to 'pending'.
    """

    def __init__(self, step_num: int, description: str, status: str = "pending"):
        self.step_num = step_num
        self.description = description
        self.status = status
        self.result: Optional[str] = None

    def __repr__(self) -> str:
        return f"Step {self.step_num}: {self.description} [{self.status}]"


class PlanningPattern(BasePattern):
    """
    Planning pattern for decomposing complex tasks into executable steps.

    This pattern implements a plan-and-execute approach where the agent first
    generates a structured plan breaking down the task into discrete steps,
    then executes each step sequentially, and finally synthesizes the results.

    The planning pattern is ideal for:
    - Complex multi-step tasks requiring coordination
    - Tasks benefiting from upfront decomposition and strategy
    - Long-running workflows where progress tracking is important
    - Tasks where step dependencies need explicit management

    The pattern follows this workflow:
    1. **Plan Generation**: Agent creates a structured plan with numbered steps
    2. **Sequential Execution**: Each step is executed in order
    3. **Result Synthesis**: Final answer integrates all step results

    This approach provides transparency, debuggability, and the ability to
    resume or modify execution mid-workflow.
    """

    def __init__(self, agent, config: Optional[PatternConfig] = None):
        """
        Initialize the Planning pattern.

        :param agent: BaseAgent - The agent to use for plan generation and step execution.
        :param config: Optional[PatternConfig] - Pattern configuration. If None, uses default with max_iterations=20.
        """
        if config is None:
            config = PatternConfig(
                name="planning", pattern_type=PatternType.PLANNING, max_iterations=20
            )
        super().__init__(config)
        self.agent = agent
        self._conversation = Conversation(id=f"planning-{id(self)}")
        self._state = State()
        self._plan: List[PlanStep] = []

    async def execute(self, task: str, context: Optional[Dict[str, Any]] = None) -> PatternResult:
        """
        Execute the planning pattern on the given task.

        The method first generates a plan by decomposing the task, then executes
        each step sequentially, and finally synthesizes all step results into a
        coherent final answer.

        :param task: str - The complex task to decompose and execute.
        :param context: Optional[Dict[str, Any]] - Additional context for execution. Defaults to None.
        :return: PatternResult - Contains success status, final answer, plan steps, and execution metadata.
        :raises PatternExecutionError: If the planning pattern execution fails.

        Example:
            >>> pattern = PlanningPattern(agent=my_agent)
            >>> result = await pattern.execute("Research and summarize the history of AI from 1950 to 2020")
            >>> print(f"Answer: {result.result}")
            >>> print(f"Plan steps: {result.metadata['plan']}")
        """
        self.reset()
        self._plan.clear()
        self._state.update_status(StateStatus.RUNNING)

        logger.info(f"Starting planning pattern for task: {task[:50]}...")

        try:
            plan = await self._create_plan(task)
            self._plan = plan
            logger.info(f"Created plan with {len(plan)} steps")

            results = []
            for step in plan:
                if self.current_iteration >= self.config.max_iterations:
                    logger.warning("Reached max iterations during plan execution")
                    break

                self.increment_iteration()
                logger.debug(f"Executing step {step.step_num}: {step.description}")

                state_data = {
                    "iteration": self.current_iteration,
                    "current_step": step.step_num,
                    "step_description": step.description,
                    "previous_results": results,
                }

                step_result = await self.step(state_data)
                step.result = step_result.get("result")
                step.status = "completed" if step_result.get("success") else "failed"
                results.append(step.result)

                if not step_result.get("success"):
                    logger.error(f"Step {step.step_num} failed: {step.result}")
                    break

            final_result = await self._synthesize_results(task, results)

            self._state.update_status(StateStatus.COMPLETED)

            return PatternResult(
                success=True,
                result=final_result,
                iterations=self.current_iteration,
                metadata={
                    "plan": [str(s) for s in self._plan],
                    "step_results": results,
                    "completed_steps": sum(1 for s in self._plan if s.status == "completed"),
                },
            )

        except Exception as e:
            logger.error(f"Planning pattern failed: {str(e)}")
            self._state.update_status(StateStatus.FAILED)
            return PatternResult(
                success=False, result=None, iterations=self.current_iteration, error=str(e)
            )

    async def step(self, state: Dict[str, Any]) -> Dict[str, Any]:
        step_num = state.get("current_step")
        description = state.get("step_description")
        previous_results = state.get("previous_results", [])

        context_str = ""
        if previous_results:
            context_str = "\n\nPrevious results:\n"
            for i, res in enumerate(previous_results, 1):
                context_str += f"Step {i}: {res}\n"

        prompt = f"Execute the following step:\n{description}{context_str}"

        try:
            result = await self.agent.generate(prompt)
            self._conversation.add_message(
                Message(role=MessageRole.ASSISTANT, content=result, metadata={"step": step_num})
            )
            return {"success": True, "result": result}
        except Exception as e:
            return {"success": False, "result": str(e)}

    async def _create_plan(self, task: str) -> List[PlanStep]:
        planning_prompt = (
            f"Create a detailed step-by-step plan to accomplish this task:\n{task}\n\n"
            "Provide your plan as a numbered list. Each step should be clear and actionable.\n"
            "Format:\n1. <step description>\n2. <step description>\netc."
        )

        self._conversation.add_message(Message(role=MessageRole.USER, content=planning_prompt))
        plan_response = await self.agent.generate(planning_prompt)
        self._conversation.add_message(Message(role=MessageRole.ASSISTANT, content=plan_response))

        plan = self._parse_plan(plan_response)
        return plan

    def _parse_plan(self, plan_text: str) -> List[PlanStep]:
        steps = []
        lines = plan_text.strip().split("\n")

        for line in lines:
            line = line.strip()
            if not line:
                continue

            if line[0].isdigit() and ("." in line or ")" in line):
                parts = line.split(".", 1) if "." in line else line.split(")", 1)
                if len(parts) == 2:
                    try:
                        step_num = int(parts[0].strip())
                        description = parts[1].strip()
                        steps.append(PlanStep(step_num, description))
                    except ValueError:
                        continue

        if not steps:
            steps.append(PlanStep(1, plan_text))

        return steps

    async def _synthesize_results(self, task: str, results: List[str]) -> str:
        synthesis_prompt = f"Original task: {task}\n\nResults from each step:\n"
        for i, result in enumerate(results, 1):
            synthesis_prompt += f"\nStep {i}: {result}\n"

        synthesis_prompt += "\nProvide a final synthesized answer based on all the steps:"

        final_answer = await self.agent.generate(synthesis_prompt)
        return final_answer

    def get_plan(self) -> List[PlanStep]:
        return self._plan.copy()

    def get_conversation(self) -> Conversation:
        return self._conversation
