from langchain_openai import ChatOpenAI
from langchain.schema import SystemMessage, HumanMessage
from langgraph.prebuilt import create_react_agent
from tools import (
    wikipedia_search_tool, arxiv_search_tool,
    audio_transcriber_tool, excel_tool, analyze_code_tool, image_tool,
    add_tool, subtract_tool, multiply_tool, divide_tool, web_search_tool
)

# ──────────────────────────────── Config ────────────────────────────────
SYSTEM_PROMPT = """
You are a helpful AI assistant. You will answer questions by thinking step-by-step, taking actions using tools when necessary, and finishing with a final answer.




Action Input: [input]

When you receive an observation, continue reasoning. Write:

IMPORTANT:
- 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 comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
- If using tools that require a `task_id`, only use the value provided.
- Do not make up tool names.
- Do not loop unnecessarily.
- Provide FINAL ANSWER fast.
- Try to use the web search tool first before using wiki and arxiv  tools.
- If wiki tool doesnt give you enough information, try again with a shorter query, or a more genral query. Dont use a query that is too specific.
- DOnt user wiki tool, web search tool and arxiv tool more than 3 times.

Return the final answer in the following format:
Thought: [final reasoning]
FINAL ANSWER: [your answer]
"""

TOOLS = [
    wikipedia_search_tool, arxiv_search_tool,
    audio_transcriber_tool, excel_tool, analyze_code_tool, image_tool,
    add_tool, subtract_tool, multiply_tool, divide_tool, web_search_tool
]

# ───────────────────────────── Agent Wrapper ─────────────────────────────
class SimpleLangGraphAgent:
    def __init__(self, model_name="gpt-4.1-mini"):
        self.agent = create_react_agent(
            model=ChatOpenAI(model_name=model_name, temperature=0.3),
            tools=TOOLS,
        )

    def run(self, question: str, task_id: str = None, max_steps: int = 10) -> str:
        messages = [SystemMessage(content=SYSTEM_PROMPT)]
        if task_id:
            messages[0].content += f"\n\nYour task_id is: {task_id}. Use it only with the tools that require it."
        messages.append(HumanMessage(content=question))

        state = {"messages": messages}
        final_state = self.agent.invoke(state)

        for msg in final_state["messages"][::-1]:
            if "FINAL ANSWER:" in msg.content.upper():
                print(f"\n\n thoughts = {msg.content.split('Thought:')[-1].strip()} \n\n")
                return msg.content.split("FINAL ANSWER:")[-1].strip()

        return "No FINAL ANSWER found."