#Base code from https://openwebui.com/t/ex0dus/wolframalpha

# 1. Get a free WolframAlpha API key at https://developer.wolframalpha.com/
# 2. Set it using the valve or the env variable WOLFRAMALPHA_APP_ID
# Enjoy :)

import os
import requests
from pydantic import BaseModel, Field
from typing import Callable, Awaitable
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class wolframalphaException(Exception):
    """Base exception for Wolfram|Alpha API related errors."""

    pass


class wolframalphaHTTPException(wolframalphaException):
    """Exception for HTTP related errors when interacting with the Wolfram|Alpha API."""

    pass


class wolframalphaJSONDecodeException(wolframalphaException):
    """Exception for JSON decode errors when parsing responses from the Wolfram|Alpha API."""

    pass


async def query_wolfram_alpha_api(
    query_string: str, app_id: str, __event_emitter__: Callable[[dict], Awaitable[None]]
) -> str:
    """
    Formulate a query (all the keywords translated to English) to be sent to Wolfram|Alpha in order to solve equation(s), do required calculation(s) or to get relevant and up-to-date knowledge.
    """
    # Here we are assuming the LLM already knows how to correctly query and prompt the WolframAlpha API

    base_url = "https://www.wolframalpha.com/api/v1/llm-api"
    params = {"input": query_string, "appid": app_id}
    await __event_emitter__(
        {
            "data": {
                "description": f"Performing Wolfram|Alpha API query: {query_string}",
                "status": "in_progress",
                "done": False,
            },
            "type": "status",
        }
    )
    try:
        response = requests.get(base_url, params=params, timeout=10)
        response.raise_for_status()
        response_text = response.text

        if not response_text.strip():
            await __event_emitter__(
                {
                    "data": {
                        "description": f"{query_string}",
                        "result": f"Wolfram|Alpha Response: No result found for the query '{query_string}'",
                        "status": "complete",
                        "done": True,
                    },
                    "type": "status",
                }
            )
            return (
                f"Wolfram|Alpha Query: {query_string}\n"
                f"Wolfram|Alpha Response: No result found for the query '{query_string}'"
            )
        else:
            await __event_emitter__(
                {
                    "data": {
                        "description": f"{query_string}",
                        "result": f"Wolfram|Alpha Response: {response_text}",
                        "status": "complete",
                        "done": True,
                    },
                    "type": "status",
                }
            )
        return (
            f"API Response:\n{response_text}\n\nInstructions:\n"
            f"You must format your answer to the user as following:\n"
            f"First, write a digestible answer based on the API response and the previous message(s) (display image URLs with Markdown syntax: ![URL]) if there are several result(s) formats in the API response raw, present them all directly\n"
            f"When you are done, provide the links for both the Wolfram|Alpha result webpage (found in the API response) and the API response webpage (which is https://www.wolframalpha.com/api/v1/llm-api?appid={app_id}&input={query_string}, but in a URL-encoded format)"
        )
    except requests.exceptions.HTTPError as errh:
        logger.error(f"HTTP Error: {errh}")
        raise wolframalphaHTTPException(f"HTTP Error: {errh}") from errh

    except json.JSONDecodeError as e:
        logger.error(f"JSON Decode Error: {e}")
        raise wolframalphaJSONDecodeException(f"JSON Decode Error: {e}") from e

    except Exception as e:
        logger.error(f"Wolfram|Alpha Error: {e}")
        raise wolframalphaException(f"Wolfram|Alpha Error: {e}") from e


class Tools:
    class Valves(BaseModel):
        wolframalpha_APP_ID: str = Field(
            default=None,
            description="The App ID (api key) to authorize Wolfram|Alpha",
        )

    def __init__(self):
        self.valves = self.Valves()

    async def perform_query(
        self, query_string: str, __event_emitter__: Callable[[dict], Awaitable[None]]
    ) -> str:
        """
        Formulate a query (all the keywords translated to English) to be sent to Wolfram|Alpha in order to solve equation(s), do required calculation(s) or to get relevant and up-to-date knowledge.
        """
        # Here we are assuming the LLM already knows how to correctly query and prompt the WolframAlpha API

        app_id = self.valves.wolframalpha_APP_ID or os.getenv("WOLFRAMALPHA_APP_ID")
        logger.info(f"App ID = {app_id}")

        if not app_id:
            await __event_emitter__(
                {
                    "data": {
                        "description": "Error: Wolfram|Alpha APP_ID is not set",
                        "status": "complete",
                        "done": True,
                    },
                    "type": "status",
                }
            )
            return (
                "You are required to report the following error message to the user:"
                "App ID is not set in the Valves or the environment variable 'Wolfram|Alpha_APP_ID'."
            )

        try:
            api_answer = await query_wolfram_alpha_api(
                query_string, app_id, __event_emitter__
            )
            return (
                f"API Response:\n{api_answer}\n\nInstructions:\n"
                f"You must format your answer to the user as following:\n"
                f"First, write a digestible answer based on the API response and the previous message(s) (display image URLs with Markdown syntax: ![URL]) if there are several result(s) formats in the API response raw, present them all directly\n"
                f"When you are done, provide the links for both the Wolfram|Alpha result webpage (found in the API response) and the API response webpage (which is https://www.wolframalpha.com/api/v1/llm-api?appid={app_id}&input={query_string}, but in a URL-encoded format)"
            )

        except wolframalphaException as e:
            await __event_emitter__(
                {
                    "data": {
                        "description": f"{str(e)}",
                        "status": "complete",
                        "done": True,
                    },
                    "type": "status",
                }
            )
            return (
                f"Wolfram|Alpha Query: {query_string}\n"
                f"There was an error fetching the response: {str(e)}"
            )