import os
import requests
from langchain_mcp_adapters.tool import MCPTool

gaia_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.
"""

class BasicAgent:
    def __init__(self, model_name="deepseek-ai/deepseek-v3.1"):
        print("BasicAgent initialized.")

        # MCP tool configuration
        self.mcp_tool = MCPTool(
            tool_name="generic_tool",  # replace with actual tool name exposed by MCP
            server_url="http://localhost:8080",
        )

        # NVIDIA NIM API configuration
        self.model_name = model_name
        self.api_key = os.getenv("NVIDIA_API_KEY")
        self.api_base = "https://integrate.api.nvidia.com/v1"

    def call_nim_api(self, user_input: str) -> str:
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": self.model_name,
            "messages": [
                {"role": "system", "content": gaia_system_prompt},
                {"role": "user", "content": user_input}
            ],
            "temperature": 0.7
        }
        response = requests.post(
            f"{self.api_base}/chat/completions",
            headers=headers,
            json=payload
        )
        try:
            return response.json()["choices"][0]["message"]["content"]
        except Exception as e:
            print("Error calling NIM API:", e)
            return "NIM API call failed."

    def __call__(self, question: str) -> str:
        print(f"Agent received input (first 50 chars): {question[:50]}...")

        # Call NVIDIA NIM API
        nim_output = self.call_nim_api(question)
        print(f"NIM response: {nim_output[:100]}...")

        # Optionally use MCP tool based on input
        if "scrape" in question.lower():
            mcp_result = self.mcp_tool.run({
                "url": "https://example.com",
                "selectors": {
                    "title": ".title",
                    "price": ".price"
                }
            })
            print("MCP result:", mcp_result)
            return f"NIM: {nim_output}\n\nMCP: {mcp_result}"

        return nim_output