from termcolor import colored
import os
from dotenv import load_dotenv
import requests
import json
import operator
import re

load_dotenv()

class OllamaModel:
    def __init__(self, model, system_prompt, temperature=0, stop=None):
        """
        用给定的参数初始化 OllamaModel。

        参数:
        model (str): 要使用的模型名称。
        system_prompt (str): 要使用的系统提示。
        temperature (float): 模型的温度设置。
        stop (str): 模型的停止标记。
        """
        self.model_endpoint = os.getenv(
            "OLLAMA_ENDPOINT",
            "http://localhost:11434/api/generate",
        )
        self.temperature = temperature
        self.model = model
        self.system_prompt = system_prompt
        self.headers = {"Content-Type": "application/json"}
        self.stop = stop

    def generate_text(self, prompt):
        """
        将用户输入和系统提示发送给本地 Ollama 服务，并返回模型生成的文本。

        参数:
        prompt (str): 用户输入。

        返回:
        str: 模型返回的文本内容。
        """
        payload = {
            "model": self.model,
            "prompt": f"{self.system_prompt}\n\n用户输入: {prompt}\n\n请按要求输出:",
            "temperature": self.temperature,
            "stream": False,
        }

        if self.stop:
            payload["stop"] = [self.stop] if isinstance(self.stop, str) else self.stop

        response = requests.post(
            self.model_endpoint,
            headers=self.headers,
            data=json.dumps(payload),
            timeout=60,
        )
        response.raise_for_status()

        result = response.json()
        return result.get("response", "").strip()

def basic_calculator(input_str):
    """
    根据输入字符串或字典对两个数字执行数值运算。

    参数:
    input_str (str或dict): 要么是表示包含'num1'、'num2'和'operation'键的字典的JSON字符串，
    要么是直接的字典。例如: '{"num1": 5, "num2": 3, "operation": "add"}'
    或{"num1": 67869, "num2": 9030393, "operation": "divide"}

    返回:
    str: 运算的格式化结果。

    抛出:
    Exception: 如果在运算过程中发生错误（例如，除以零）。
    ValueError: 如果请求了不支持的运算或输入无效。
    """

    try:
        if isinstance(input_str, dict):
            input_dict = input_str
        else:
            input_str_clean = input_str.replace("'", "\"")
            input_str_clean = input_str_clean.strip().strip("\"")
            input_dict = json.loads(input_str_clean)
        if not all(key in input_dict for key in ['num1', 'num2', 'operation']):
            return "Error: Input must contain 'num1', 'num2', and 'operation'."
        num1 = float(input_dict['num1'])
        num2 = float(input_dict['num2'])
        operation = input_dict['operation'].lower()
    except (json.JSONDecodeError, KeyError) as e:
        return "Invalid input format. Please provide valid numbers and operation."
    except ValueError as e:
        return "Error: Please provide valid numerical values."

    operations = {
        'add': operator.add,
        'plus': operator.add,
        'subtract': operator.sub,
        'minus': operator.sub,
        'multiply': operator.mul,
        'times': operator.mul,
        'divide': operator.truediv,
        'floor_divide': operator.floordiv,
        'mod': operator.mod,
        'power': operator.pow,
    }

    if operation not in operations:
        return f"Error: Unsupported operation '{operation}'."

    if operation in ['divide', 'floor_divide', 'mod'] and num2 == 0:
        return "Error: Division by zero is not allowed."

    try:
        result = operations[operation](num1, num2)
        return f"Result: {result}"
    except Exception as e:
        return f"Error during calculation: {str(e)}"


def reverse_string(input_str):
    """
    反转输入文本。

    参数:
    input_str (str或dict): 要反转的文本，或包含 text/input/tool_input 字段的字典。

    返回:
    str: 反转后的文本。
    """
    if isinstance(input_str, dict):
        text = (
            input_str.get("text")
            or input_str.get("input")
            or input_str.get("tool_input")
            or ""
        )
    else:
        text = str(input_str)

    text = text.strip()

    # 兼容模型把完整请求当作 tool_input 返回的情况。
    text = re.sub(r"^(can you\s+)?reverse( the word)?( of)?\s+", "", text, flags=re.I)
    text = text.strip().strip("?").strip()
    text = text.strip("'\"")

    return text[::-1]

class ToolBox:
    def __init__(self):
        self.tools_dict = {}

    def store(self, functions_list):
        """
        存储列表中每个函数的名称和文档字符串。
        参数:
        functions_list (list): 要存储的函数对象列表。
        返回:
        dict: 以函数名称为键，其文档字符串为值的字典。
        """
        for func in functions_list:
            self.tools_dict[func.__name__] = func.__doc__
        return self.tools_dict

    def tools(self):
        """
        将store方法中创建的字典转换为文本字符串返回。
        返回:
        str: 存储的函数及其文档字符串的字典，以文本字符串形式返回。
        """
        tools_str = ""
        for name, doc in self.tools_dict.items():
            tools_str += f"{name}: \"{doc}\"\n"
        return tools_str.strip()

AGENT_SYSTEM_PROMPT = """
你是一个工具路由型 AI Agent。你的任务是理解用户请求，决定是否需要调用工具，并输出严格的 JSON。

可用工具:

1. basic_calculator: 用于数学运算
   - 示例输入: "Calculate 15 plus 7"
   - 示例输出: {"tool_choice": "basic_calculator", "tool_input": {"num1": 15, "num2": 7, "operation": "plus"}}
   - 示例输入: "What is 100 divided by 5?"
   - 示例输出: {"tool_choice": "basic_calculator", "tool_input": {"num1": 100, "num2": 5, "operation": "divide"}}

2. reverse_string: 用于文本反转
   - 示例输入: "Reverse of 'Howwww'?"
   - 示例输出: {"tool_choice": "reverse_string", "tool_input": "Howwww"}
   - 示例输入: "What is the reverse of Python?"
   - 示例输出: {"tool_choice": "reverse_string", "tool_input": "Python"}

3. no tool: 用于一般对话和问题
   - 示例输入: "Who are you?"
   - 示例输出: {"tool_choice": "no tool", "tool_input": "I am an AI assistant that can answer questions, perform calculations, and reverse text."}
   - 示例输入: "How are you?"
   - 示例输出: {"tool_choice": "no tool", "tool_input": "I'm functioning well and ready to help."}

严格规则:

1. 对于关于身份、能力或感受的问题:
   - 始终使用 "no tool"
   - 提供完整、友好的回复
   - 提及你的能力

2. 对于任何文本反转请求:
   - 始终使用 "reverse_string"
   - 仅提取要反转的文本
   - 去除引号、"reverse of" 和其他多余文本

3. 对于任何数学运算:
   - 始终使用 "basic_calculator"
   - 提取数字和运算
   - 将文本形式的数字转换为数字

下面是你的工具列表及其描述:
{tool_descriptions}

注意:
- 你的回复必须始终是包含 "tool_choice" 和 "tool_input" 字段的有效 JSON。
- 不要输出 Markdown。
- 不要输出解释文字。
"""


class Agent:
    def __init__(self, tools, model_service, model_name, stop=None):
        """
        用工具列表和模型初始化智能体。

        参数:
        tools (list): 可调用工具函数列表。
        model_service (class): 模型服务类，例如 OllamaModel。
        model_name (str): 要使用的模型名称。
        stop (str): 模型停止标记。
        """
        self.tools = tools
        self.tools_map = {tool.__name__: tool for tool in tools}
        self.tool_box = ToolBox()
        self.tool_descriptions = self.prepare_tools()
        self.system_prompt = AGENT_SYSTEM_PROMPT.replace(
            "{tool_descriptions}",
            self.tool_descriptions,
        )
        self.model = model_service(
            model=model_name,
            system_prompt=self.system_prompt,
            temperature=0,
            stop=stop,
        )

    def prepare_tools(self):
        """
        存储工具函数，并返回工具描述文本。
        """
        self.tool_box.store(self.tools)
        return self.tool_box.tools()

    def _extract_json(self, text):
        """
        从模型输出中提取 JSON 对象。

        有些模型会额外输出说明或代码围栏，因此这里做一层容错解析。
        """
        text = text.strip()
        text = text.replace("```json", "").replace("```", "").strip()

        try:
            return json.loads(text)
        except json.JSONDecodeError:
            match = re.search(r"\{.*\}", text, flags=re.S)
            if not match:
                raise ValueError(f"Model did not return JSON: {text}")
            return json.loads(match.group(0))

    def think(self, prompt):
        """
        调用模型，让模型决定工具选择与工具参数。

        参数:
        prompt (str): 用户输入。

        返回:
        dict: 包含 tool_choice 和 tool_input 的决策结果。
        """
        response = self.model.generate_text(prompt)
        decision = self._extract_json(response)

        if "tool_choice" not in decision or "tool_input" not in decision:
            raise ValueError(f"Invalid tool decision: {decision}")

        return decision

    def work(self, prompt):
        """
        执行一次完整 Agent 工作流：思考、选择工具、执行工具、输出结果。

        参数:
        prompt (str): 用户输入。

        返回:
        str: 最终输出。
        """
        try:
            decision = self.think(prompt)
            tool_choice = decision["tool_choice"]
            tool_input = decision["tool_input"]

            print(colored(f"Tool choice: {tool_choice}", "cyan"))
            print(colored(f"Tool input: {tool_input}", "cyan"))

            if tool_choice == "no tool":
                final_answer = str(tool_input)
            else:
                tool = self.tools_map.get(tool_choice)
                if not tool:
                    final_answer = f"Error: Unknown tool '{tool_choice}'."
                else:
                    final_answer = tool(tool_input)

            print(colored(f"Final answer: {final_answer}", "green"))
            return final_answer
        except Exception as e:
            error_message = f"Agent error: {str(e)}"
            print(colored(error_message, "red"))
            return error_message

if __name__ == "__main__":
    """
    使用此智能体的说明:
    你可以尝试的示例查询:
    1. 计算器运算:
        - "Calculate 15 plus 7"
        - "What is 100 divided by 5?"
        - "Multiply 23 and 4"
    2. 字符串反转:
        - "Reverse the word 'hello world'"
        - "Can you reverse 'Python Programming'?"
    3. 一般问题（将得到直接回复）:
        - "Who are you?"
        - "What can you help me with?"
    Ollama命令（在终端中运行这些命令）:
        - 查看可用模型:      'ollama list'
        - 查看正在运行的模型: 'ps aux | grep ollama'
        - 列出模型标签:      'curl http://localhost:11434/api/tags'
        - 拉取新模型:        'ollama pull mistral'
        - 运行模型服务器:    'ollama serve'
    """
    tools = [basic_calculator, reverse_string]
    model_service = OllamaModel
    model_name = "llama2"
    stop = "<<|eot_id|>"
    agent = Agent(
        tools=tools,
        model_service=model_service,
        model_name=model_name,
        stop=stop,
    )

    print("\nWelcome to the AI Agent! Type 'exit' to quit.")
    print("You can ask me to:")
    print("1. Perform calculations (e.g., 'Calculate 15 plus 7')")
    print("2. Reverse strings (e.g., 'Reverse hello world')")
    print("3. Answer general questions\n")

    while True:
        prompt = input("Ask me anything: ")
        if prompt.lower() == "exit":
            break
        agent.work(prompt)