from typing import Dict, Any, List
from loguru import logger
from utils.llm_orchestrator import LLMOrchestrator


class PlanningAgent:
    def __init__(self, llm_api_key: str):
        """Initialize the Planning Agent."""
        logger.info("Initializing PlanningAgent")
        self.llm_orchestrator = LLMOrchestrator(llm_api_key)
        self.capabilities = [
            "task_planning",
            "goal_decomposition",
            "plan_refinement",
            "task_prioritization"
        ]
        self.setup_logger()

    def setup_logger(self):
        """Configure logging for the agent."""
        logger.add("logs/planning_agent.log", rotation="500 MB")

    async def generate_plan(
            self, goal: str, available_agents: List[str]) -> Dict[str, Any]:
        """Generate a task plan based on a high-level goal."""
        logger.info(f"Generating plan for goal: {goal}")
        try:
            prompt = f"""
            You are an expert planner. Generate a detailed task plan to achieve the following goal:
            Goal: {goal}

            Available agents: {', '.join(available_agents)}

            Think step-by-step and explain your reasoning for each step.
            The plan should be a list of steps, each with:
            - A clear description of the task.
            - The agent best suited to execute the task.
            - Any necessary input or parameters for the task.

            Example:
            1. Task: Summarize the latest news on topic X.
               Agent: web_browsing_agent
               Input: topic=X
               Reasoning: To get the latest news, we need to use the web_browsing_agent to search for news on topic X.
            2. Task: Analyze the sentiment of the news summary.
               Agent: data_analysis_agent
               Input: summary from step 1
               Reasoning: To analyze the sentiment, we can use the data_analysis_agent to process the summary from the previous step.
            """
            plan_str = await self.llm_orchestrator.generate_completion(prompt)
            plan = self.parse_plan(plan_str)
            logger.info(f"Plan generated successfully: {plan}")
            return {
                "status": "success",
                "plan": plan
            }

        except Exception as e:
            logger.error(f"Error generating plan: {str(e)}")
            return {
                "status": "error",
                "message": str(e)
            }

    def parse_plan(self, plan_str: str) -> List[Dict[str, Any]]:
        """Parse the plan generated by the LLM into a structured format."""
        plan = []
        steps = plan_str.strip().split("\n")
        current_step = {}

        for line in steps:
            if line.startswith(tuple(f"{i}." for i in range(1, 10))):
                if current_step:
                    plan.append(current_step)
                current_step = {"task": line.split("Task: ")[1]}
            elif "Agent: " in line:
                current_step["agent"] = line.split("Agent: ")[1]
            elif "Input: " in line:
                current_step["input"] = line.split("Input: ")[1]
            elif "Reasoning: " in line:
                current_step["reasoning"] = line.split("Reasoning: ")[1]

        if current_step:
            plan.append(current_step)

        return plan