"""
Orchestrator agent using ReAct pattern for complex multi-step queries.
Coordinates multiple specialist agents to answer comprehensive questions.
"""
from typing import List
from langchain.agents import AgentExecutor, create_react_agent
from langchain.tools import Tool
from langchain_core.prompts import PromptTemplate


from agentic_student_assistant.core.base.base_agent import BaseAgent
from agentic_student_assistant.core.utils.config_loader import get_config, get_prompt


class OrchestratorAgent(BaseAgent):
    """
    Orchestrator agent that uses ReAct to coordinate multiple specialist agents.
    Handles complex queries requiring multiple steps or multiple domain expertise.
    """
    
    def __init__(self):
        """Initialize orchestrator with ReAct agent and tools."""
        config = get_config()
        super().__init__(config, agent_name="orchestrator")
        
        # Create tools for each specialist agent
        self.tools = self._create_tools()
        
        # Create ReAct agent
        self.agent = self._create_react_agent()
        self.agent_executor = AgentExecutor(
            agent=self.agent,
            tools=self.tools,
            verbose=True,
            max_iterations=5,
            handle_parsing_errors=True
        )
    
    def _create_tools(self) -> List[Tool]:
        """
        Create tools for each specialist agent.
        
        Returns:
            List of LangChain tools
        """
        from agentic_student_assistant.talk2jobs.agents.job_market_agent import JobMarketAgent # pylint: disable=import-outside-toplevel
        from agentic_student_assistant.talk2books.agents.books_recommend_agent import BooksRecommendAgent # pylint: disable=import-outside-toplevel
        from agentic_student_assistant.talk2papers.agents.paper_recommend_agent import PaperRecommendAgent # pylint: disable=import-outside-toplevel
        
        # Initialize agents (lazy loading to avoid circular imports)
        job_agent = None
        books_agent = None
        papers_agent = None
        
        def get_job_agent():
            nonlocal job_agent
            if job_agent is None:
                job_agent = JobMarketAgent()
            return job_agent
        
        def get_books_agent():
            nonlocal books_agent
            if books_agent is None:
                books_agent = BooksRecommendAgent()
            return books_agent
        
        def get_papers_agent():
            nonlocal papers_agent
            if papers_agent is None:
                papers_agent = PaperRecommendAgent()
            return papers_agent
        
        tools = [
            Tool(
                name="JobMarketSearch",
                func=lambda q: get_job_agent().process(q),
                description=(
                    "Search job listings and career opportunities. Use this to find "
                    "current job openings, understand job market demand, or get "
                    "information about specific roles and their requirements."
                )
            ),
            Tool(
                name="BookRecommendations",
                func=lambda q: get_books_agent().process(q),
                description=(
                    "Find academic book recommendations and learning resources. Use this "
                    "to suggest textbooks, monographs, or other high-quality learning "
                    "materials for specific topics."
                )
            ),
            Tool(
                name="PaperRecommendations",
                func=lambda q: get_papers_agent().process(q),
                description=(
                    "Find scientific research papers and academic articles. Use this to "
                    "find primary sources, latest research, citations, and technical "
                    "details from ArXiv, Semantic Scholar, and CORE."
                )
            )
        ]
        
        return tools
    
    def _create_react_agent(self):
        """
        Create ReAct agent with appropriate prompts.
        
        Returns:
            ReAct agent
        """
        # Get prompts from config
        system_prompt = get_prompt("orchestrator_system")
        react_template = get_prompt("react_template")
        
        # Format the template with system prompt safely using replace
        # to avoid KeyError for {tools}, {input}, etc.
        formatted_template = react_template.replace("{system_prompt}", system_prompt)
        
        prompt = PromptTemplate.from_template(formatted_template)
        
        return create_react_agent(
            llm=self.llm,
            tools=self.tools,
            prompt=prompt
        )
    
    def process(self, query: str, **kwargs) -> str:
        """
        Process complex query using ReAct orchestration.
        
        Args:
            query: Complex user query
            **kwargs: Additional parameters
            
        Returns:
            Comprehensive answer from orchestration
        """
        try:
            orch_result = self.agent_executor.invoke({"input": query})
            return orch_result.get("output", "Unable to process query")
        except Exception as e: # pylint: disable=broad-exception-caught
            print(f"❌ Orchestration error: {e}")
            return (
                f"I encountered an error while processing your complex query: {str(e)}. "
                "Please try breaking it down into simpler questions."
            )


if __name__ == "__main__":
    print("🎭 Orchestrator Agent Test\n")
    
    agent = OrchestratorAgent()
    
    # Test complex query
    test_query = "What courses should I take to get an AI job in Berlin?"
    print(f"Query: {test_query}\n")
    
    result = agent.process(test_query)
    print(f"\n✅ Final Answer:\n{result}")
