from typing import List, Dict, Any
from langchain.agents import create_openai_functions_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.tools import BaseTool
from langchain_openai import ChatOpenAI

class ResearchAgent:
    def __init__(self, tools: List[BaseTool], openai_api_key: str):
        self.tools = tools
        self.llm = ChatOpenAI(
            temperature=0,
            model="gpt-4-turbo-preview",
            openai_api_key=openai_api_key
        )
        
        # Define the system prompt
        system_prompt = """You are a specialized research assistant focused on scientific literature analysis.
        Your goal is to help users find, analyze, and understand scientific papers and research findings.
        You have access to tools that can:
        1. Search for relevant papers and research
        2. Analyze PDF documents
        3. Track citations and research impact
        
        Always be thorough in your analysis and provide clear, well-structured responses.
        If you're unsure about something, be honest and ask for clarification."""
        
        # Create the prompt template
        prompt = ChatPromptTemplate.from_messages([
            ("system", system_prompt),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])
        
        # Create the agent
        self.agent = create_openai_functions_agent(
            llm=self.llm,
            prompt=prompt,
            tools=self.tools
        )
        
        # Create the agent executor
        self.agent_executor = AgentExecutor(
            agent=self.agent,
            tools=self.tools,
            verbose=False
        )
    
    def run(self, query: str, chat_history: List[Dict[str, Any]] = None) -> str:
        """Run the agent with the given query and chat history."""
        if chat_history is None:
            chat_history = []
            
        result = self.agent_executor.invoke({
            "input": query,
            "chat_history": chat_history
        })
        
        return result["output"]