import re
import constants
import time
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.agents import AgentExecutor, create_tool_calling_agent, create_openai_functions_agent
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import SystemMessage

# --- Custom Tools ---
from internet_search_tool import internet_search
from nb_tool import get_team_players_by_season
from nb_tool import get_npb_player_info

PROMPT = """
    You are an agent specialized in answering questions related to Nippon Professional Baseball players. Report your thoughts, and 
    finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. 
"""

class NpbAgent:
    def __init__(self):
        llm = ChatGoogleGenerativeAI(
            model=constants.MODEL,
            api_key=constants.API_KEY, 
            temperature=0.4,
            timeout=20)

        tools = [
            get_npb_player_info,
            get_team_players_by_season
        ]

        prompt = ChatPromptTemplate.from_messages([
            SystemMessage(content=PROMPT),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])

        agent = create_tool_calling_agent(llm, tools, prompt=prompt)
        self.executor = AgentExecutor(
            agent=agent, 
            tools=tools, 
            verbose=True,
            max_iterations=15)

    def __call__(self, question: str) -> str:
        print(f"NpbAgent agent received: {question[:50]}...")
        result = self.executor.invoke({
            "input": question,
            "chat_history": []
        })
        output = result.get("output", "No answer returned.")
        print(f"Agent response: {output}")
        match = re.search(r"FINAL ANSWER:\s*(.*)", output)
        if match:
            return match.group(1).strip()
        else:
            return output