import os

from langgraph.prebuilt import create_react_agent
from langchain_groq import ChatGroq
from tools import all_tools

SYSTEM_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 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."
    "You have access to a web search tool, a webpage fetcher tool, an image identification tool, and a YouTube identifier tool. "
    "Use the search tool when you need to look up factual information. "
    "Use the webpage fetcher tool when you need to read the actual contents of a specific URL (e.g. a Wikipedia article, a research paper, or any link found via search). First search, then fetch the relevant URL to get detailed information. "
    "Use the image identification tool when the question involves an image URL. "
    "When the question includes a YouTube link (youtube.com or youtu.be), use the YouTube identifier tool with that URL first; it returns title, channel, and transcript text so you can describe what is said or shown in the video before answering. "
    "STRATEGY: When answering factual questions, follow these steps: "
    "1) Search for the topic. "
    "2) If the search results mention a relevant URL (e.g. a Wikipedia page), fetch that URL with the webpage fetcher to get detailed information. "
    "3) If the first search didn't give enough detail, try a more specific or different search query. "
    "Do not guess — always verify facts by reading the actual source page.\n"
    "For logic puzzles, reversed text, math problems, or categorization tasks, think carefully step by step before answering. "
    "For reversed text, reverse the string character by character to decode it.\n"
    "IMPORTANT: When calling any tool, pass ONLY the required arguments as plain values. "
    "Do NOT include your reasoning, thoughts, or explanations in tool arguments. "
    "For example, when searching, pass only the search query string like 'population of France' — never 'I need to find out the population of France'."
)

chat_model = ChatGroq(
    model="meta-llama/llama-4-scout-17b-16e-instruct",
    temperature=0,
    api_key=os.environ.get("GROQ_API_KEY"),
)

agent = create_react_agent(
    chat_model,
    all_tools,
    prompt=SYSTEM_PROMPT,
)


_MAX_RETRIES = 3


def run_agent(question: str) -> str:
    """Invoke the LangGraph agent and return only the text after 'FINAL ANSWER:'."""
    for attempt in range(_MAX_RETRIES):
        try:
            result = agent.invoke(
                {"messages": [("user", question)]},
                {"recursion_limit": 50},
            )
            messages = result["messages"]
            content = messages[-1].content
            if isinstance(content, list):
                content = "\n".join(
                    block.get("text", "") if isinstance(block, dict) else str(block)
                    for block in content
                )
            print(content)
            if "FINAL ANSWER:" in content:
                return content.split("FINAL ANSWER:")[-1].strip()
            return content
        except Exception as e:
            print(f"Attempt {attempt + 1}/{_MAX_RETRIES} failed: {e}")
            if attempt == _MAX_RETRIES - 1:
                return f"Error after {_MAX_RETRIES} attempts: {e}"