agents/planner_agent.py

import os
import logging
from typing import List, Dict

from llama_index.core.agent.workflow import ReActAgent
from llama_index.core.tools import FunctionTool
from llama_index.llms.google_genai import GoogleGenAI


# Setup logging
logger = logging.getLogger(__name__)

# Helper function to load prompt from file
def load_prompt_from_file(filename: str, default_prompt: str) -> str:
    """Loads a prompt from a text file."""
    try:
        # Assuming the prompt file is in the same directory as the agent script
        script_dir = os.path.dirname(__file__)
        prompt_path = os.path.join(script_dir, filename)
        with open(prompt_path, "r") as f:
            prompt = f.read()
            logger.info(f"Successfully loaded prompt from {prompt_path}")
            return prompt
    except FileNotFoundError:
        logger.warning(f"Prompt file {filename} not found at {prompt_path}. Using default.")
        return default_prompt
    except Exception as e:
        logger.error(f"Error loading prompt file {filename}: {e}", exc_info=True)
        return default_prompt

# --- Tool Functions ---

def plan(objective: str) -> List[str]:
    """
    Generate a list of sub-steps (4-8) from the given objective using an LLM.
    Args:
        objective (str): The research or task objective.
    Returns:
        List[str]: A list of sub-steps as strings, or an error message list.
    """
    logger.info(f"Generating plan for objective: {objective[:100]}...")
    
    # Configuration for planning LLM
    planner_llm_model = os.getenv("PLANNER_TOOL_LLM_MODEL", "gemini-2.5-pro-preview-03-25") # Specific model for this tool?
    gemini_api_key = os.getenv("GEMINI_API_KEY")
    if not gemini_api_key:
        logger.error("GEMINI_API_KEY not found for planning tool LLM.")
        return "Error: GEMINI_API_KEY not set for planning."

    # Prompt for the LLM to generate sub-steps
    input_prompt = (
        "You are a research assistant. "
        "Given an objective, break it down into a list of 4-8 concise, actionable sub-steps. "
        "Ensure the steps are logically ordered.\n"
        f"Objective: {objective}\n"
        "Sub-steps (one per line, numbered):"
    )

    try:
        llm = GoogleGenAI(api_key=gemini_api_key, model="gemini-2.5-pro-preview-03-25", temperature=0.05)
        logger.info(f"Using planning LLM: {planner_llm_model}")
        response = llm.complete(input_prompt)
        
        # Post-process: split lines into sub-steps, remove numbering if present
        lines = response.text.strip().split("\n")
        sub_steps = []
        for line in lines:
            line = line.strip()
            if not line:
                continue
            # Remove potential leading numbering (e.g., "1. ", "- ")
            if line and line[0].isdigit() and "." in line[:3]:
                text = line.split(".", 1)[1].strip()
            elif line.startswith("- "):
                 text = line[2:].strip()
            else:
                text = line
            
            if text:
                sub_steps.append(text)
                
        if not sub_steps:
             logger.warning("LLM generated no sub-steps for the objective.")
             return "Error: Failed to generate sub-steps."
             
        logger.info(f"Generated {len(sub_steps)} sub-steps.")

        return sub_steps
        
    except Exception as e:
        logger.error(f"LLM call failed during planning: {e}", exc_info=True)
        return f"Error during planning: {e}"

def synthesize_and_report(results: List[Dict[str, str]]) -> str:
    """
    Aggregate results from sub-steps into a coherent final report using an LLM.
    Args:
        results (List[Dict[str, str]]): List of dictionaries, each with "sub_step" and "answer" keys.
    Returns:
        str: A unified, well-structured report, or an error message.
    """
    logger.info(f"Synthesizing results from {len(results)} sub-steps...")
    if not results:
        logger.warning("Synthesize called with empty results list.")
        return "No results provided to synthesize."

    # Format the results for the synthesis prompt
    summary_blocks = ""
    for i, result in enumerate(results):
        sub_step = result.get("sub_step", f"Step {i+1}")
        answer = result.get("answer", "No answer provided.")
        summary_blocks += f"Sub-step {i+1}: {sub_step}\nAnswer {i+1}: {answer}\n\n"

    # Configuration for synthesis LLM
    synthesizer_llm_model = os.getenv("SYNTHESIZER_LLM_MODEL", "gemini-2.5-pro-preview-03-25") # Specific model?
    gemini_api_key = os.getenv("GEMINI_API_KEY")
    if not gemini_api_key:
        logger.error("GEMINI_API_KEY not found for synthesis tool LLM.")
        return "Error: GEMINI_API_KEY not set for synthesis."

    # Prompt for the LLM
    input_prompt = f"""You are an expert synthesizer. Given the following sub-steps and their answers derived 
    from an initial objective, produce a single, coherent, comprehensive final report that 
    addresses the original objective:

    --- SUB-STEP RESULTS ---
    {summary_blocks.strip()}
    --- END SUB-STEP RESULTS ---

    Generate the Final Report:
    """

    try:
        llm = GoogleGenAI(api_key=gemini_api_key, model="gemini-2.5-pro-preview-03-25", temperature=0.05)
        logger.info(f"Using synthesis LLM: {synthesizer_llm_model}")
        response = llm.complete(input_prompt)
        logger.info("Synthesis successful.")
        return response.text
    except Exception as e:
        logger.error(f"LLM call failed during synthesis: {e}", exc_info=True)
        return f"Error during synthesis: {e}"

def answer_question(question: str) -> str:
    """
    Answer any question by following this strict format:
      1. Include your chain of thought (your reasoning steps).
      2. End your reply with the exact template:
         FINAL ANSWER: [YOUR FINAL ANSWER]
    YOUR FINAL ANSWER must be:
      - A number, or
      - As few words as possible, or
      - A comma-separated list of numbers and/or strings.
    Formatting rules:
      * If asked for a number, do not use commas or units (e.g., $, %), unless explicitly requested.
      * If asked for a string, do not include articles or abbreviations (e.g., city names), and write digits in plain text.
      * If asked for a comma-separated list, apply the above rules to each element.
    This tool should be invoked immediately after completing the final planning sub-step.
    """
    logger.info(f"Answering question: {question[:100]}")

    gemini_api_key = os.getenv("GEMINI_API_KEY")
    if not gemini_api_key:
        logger.error("GEMINI_API_KEY not set for answer_question tool.")
        return "Error: GEMINI_API_KEY not set."

    model_name = os.getenv("ANSWER_TOOL_LLM_MODEL", "gemini-2.5-pro-preview-03-25")

    # Build the assistant prompt enforcing the required format
    assistant_prompt = (
        "You are a general AI assistant. I will ask you a question. "
        "Report your thoughts, and finish your answer with the following template: "
        "FINAL ANSWER: [YOUR FINAL ANSWER]. "
        "YOUR FINAL ANSWER should be a number OR as few words as possible "
        "OR a comma separated list of numbers and/or strings. "
        "If you are asked for a number, don't use commas for thousands or any units like $ or % unless specified. "
        "If you are asked for a string, omit articles and abbreviations, and write digits in plain text. "
        "If you are asked for a comma separated list, apply these rules to each element.\n\n"
        f"Question: {question}\n"
        "Answer:"
    )

    try:
        llm = GoogleGenAI(api_key=gemini_api_key, model="gemini-2.5-pro-preview-03-25", temperature=0.05)
        logger.info(f"Using answer LLM: {model_name}")
        response = llm.complete(assistant_prompt)
        logger.info("Answer generated successfully.")
        return response.text
    except Exception as e:
        logger.error(f"LLM call failed during answer generation: {e}", exc_info=True)
        return f"Error during answer generation: {e}"

# --- Tool Definitions ---
synthesize_tool = FunctionTool.from_defaults(
    fn=synthesize_and_report,
    name="synthesize_and_report",
    description=(
        "Aggregates results from multiple sub-steps into a final coherent report. "
        "Input: results (List[Dict[str, str]]) where each dict has \"sub_step\" and \"answer\". "
        "Output: A unified report (str) or error message."
    ),
)

generate_substeps_tool = FunctionTool.from_defaults(
    fn=plan,
    name="generate_substeps",
    description=(
        "Decomposes a high-level objective into a concise roadmap of 4–8 actionable sub-steps using an LLM. "
        "Input: objective (str). Output: List of sub-step strings (List[str]) or error list."
    )
)

answer_question = FunctionTool.from_defaults(
    fn=answer_question,
    name="answer_question",
    description=(
        "Answers any question and returns the full text, always ending with "
        "‘FINAL ANSWER: ...’ in accordance with the formatting rules."
    ),
)

# --- Agent Initialization ---
def initialize_planner_agent() -> ReActAgent:
    """Initializes the Planner Agent."""
    logger.info("Initializing PlannerAgent...")

    # Configuration for the agent's main LLM
    agent_llm_model = os.getenv("PLANNER_AGENT_LLM_MODEL", "gemini-2.5-pro-preview-03-25")
    gemini_api_key = os.getenv("GEMINI_API_KEY")

    if not gemini_api_key:
        logger.error("GEMINI_API_KEY not found for PlannerAgent.")
        raise ValueError("GEMINI_API_KEY must be set for PlannerAgent")

    try:
        llm = GoogleGenAI(api_key=gemini_api_key, model="gemini-2.5-pro-preview-03-25", temperature=0.05)
        logger.info(f"Using agent LLM: {agent_llm_model}")

        # Load system prompt
        default_system_prompt = ("You are PlannerAgent... [Default prompt content - replace with actual]" # Placeholder
                              )
        system_prompt = load_prompt_from_file("../prompts/planner_agent_prompt.txt", default_system_prompt)
        if system_prompt == default_system_prompt:
             logger.warning("Using default/fallback system prompt for PlannerAgent.")

        # Define available tools
        tools = [generate_substeps_tool, synthesize_tool]

        # Define valid handoff targets
        valid_handoffs = [
            "code_agent",
            "research_agent",
            "math_agent",
            "role_agent",
            "image_analyzer_agent",
            "text_analyzer_agent",
            "reasoning_agent",
            "long_context_management_agent",
            "advanced_validation_agent",
            "video_analyzer_agent"
        ]

        agent = ReActAgent(
            name="planner_agent",
            description=(
                "Strategically plans tasks by breaking down objectives into sub-steps using `generate_substeps`. "
                "Orchestrates execution by handing off sub-steps to specialized agents. "
                "Synthesizes final results using `synthesize_and_report`."
            ),
            tools=tools,
            llm=llm,
            system_prompt=system_prompt,
            can_handoff_to=valid_handoffs,
        )
        logger.info("PlannerAgent initialized successfully.")
        return agent

    except Exception as e:
        logger.error(f"Error during PlannerAgent initialization: {e}", exc_info=True)
        raise

# Example usage (for testing if run directly)
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    logger.info("Running planner_agent.py directly for testing...")

    # Ensure API key is set
    if not os.getenv("GEMINI_API_KEY"):
        print("Error: GEMINI_API_KEY environment variable not set. Cannot run test.")
    else:
        try:
            # Test plan generation
            print("\nTesting plan generation...")
            test_objective = "Analyze the market trends for electric vehicles in Europe for 2024."
            substeps = plan(test_objective)
            print(f"Generated Sub-steps:\n{substeps}")

            # Test synthesis
            print("\nTesting synthesis...")
            test_results = [
                {"sub_step": "Identify key EV manufacturers in Europe.", "answer": "Tesla, VW, Stellantis, Renault."}, 
                {"sub_step": "Find recent sales data.", "answer": "EV sales grew 25% year-over-year in Q1 2024."}, 
                {"sub_step": "Analyze government incentives.", "answer": "Germany reduced subsidies, France maintained them."}
            ]
            report = synthesize_and_report(test_results)
            print(f"Synthesized Report:\n{report}")

            # Initialize the agent (optional)
            # test_agent = initialize_planner_agent()
            # print("\nPlanner Agent initialized successfully for testing.")

        except Exception as e:
            print(f"Error during testing: {e}")