from runtime import Args
from typings.exec_playbook.exec_playbook import Input, Output
import requests
import json
import os
import logging
from typing import Dict, Any, Optional

"""
每个文件需要导出一个名为`handler`的函数。这是工具的入口点。

参数:
args: 入口函数的参数。
args.input - 输入参数，您可以通过args.input.xxx获取测试输入值。
args.logger - 用于打印日志的记录器实例，由运行时注入。

请记得在Metadata中填写输入/输出，这有助于LLM识别和使用工具。

返回:
函数的返回数据，应与声明的输出参数匹配。
"""

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 从环境变量获取API URL，如果未设置则使用默认值
API_URL = os.getenv("ANSIBLE_API_URL", "http://<your-ansible-api-host>:5000/run")

def handler(args: Args[Input]) -> Output:
    """
    处理函数，调用Ansible API执行playbook

    参数:
        args: 包含输入参数和日志记录器的对象

    返回:
        包含执行结果的字典（不包含stats字段）
    """
    # 获取输入参数并转换为基本类型
    playbook = args.input.playbook
    inventory = getattr(args.input, 'inventory', "localhost,")

    # 从输入参数获取API Token，如果未提供则使用默认值
    api_token = getattr(args.input, 'api_token', "<your-api-token>")

    # 记录请求信息
    logger.info(f"Executing playbook: {playbook}")
    logger.info(f"Inventory: {inventory}")

    # 准备请求头
    headers = {
        "Content-Type": "application/json",
        "X-API-Token": api_token
    }

    # 准备请求体 - 确保所有值都是基本类型
    payload = {
        "playbook": str(playbook),
        "inventory": str(inventory)
    }

    try:
        # 发送POST请求到API
        logger.info(f"Sending request to {API_URL}")
        logger.info(f"Payload: {json.dumps(payload)}")  # 记录序列化后的payload

        response = requests.post(
            API_URL,
            headers=headers,
            data=json.dumps(payload),  # 现在可以安全序列化
            timeout=60  # 设置60秒超时
        )

        # 检查响应状态码
        if response.status_code == 200:
            result = response.json()
            logger.info(f"API response: {result}")

            # 提取需要的信息，排除stats
            return {
                "status": result.get("status", "unknown"),
                "rc": result.get("rc", -1),
                "stdout": result.get("stdout", ""),
                "stderr": result.get("stderr", ""),
                "error": None
            }
        elif response.status_code == 401:
            error_msg = "认证失败，请检查API Token是否正确"
            logger.error(error_msg)
            return {
                "status": "error",
                "rc": 401,
                "stdout": "",
                "stderr": error_msg,
                "error": error_msg
            }
        elif response.status_code == 400:
            error_msg = f"请求错误: {response.json().get('error', '未知错误')}"
            logger.error(error_msg)
            return {
                "status": "error",
                "rc": 400,
                "stdout": "",
                "stderr": error_msg,
                "error": error_msg
            }
        else:
            error_msg = f"服务器返回错误: HTTP {response.status_code}"
            logger.error(error_msg)
            return {
                "status": "error",
                "rc": response.status_code,
                "stdout": "",
                "stderr": error_msg,
                "error": error_msg
            }

    except requests.exceptions.Timeout:
        error_msg = "请求超时，请检查playbook执行时间或增加超时设置"
        logger.error(error_msg)
        return {
            "status": "error",
            "rc": -1,
            "stdout": "",
            "stderr": error_msg,
            "error": error_msg
        }
    except requests.exceptions.RequestException as e:
        error_msg = f"请求失败: {str(e)}"
        logger.error(error_msg)
        return {
            "status": "error",
            "rc": -1,
            "stdout": "",
            "stderr": error_msg,
            "error": error_msg
        }
    except json.JSONDecodeError as e:
        error_msg = f"解析响应失败: {str(e)}"
        logger.error(error_msg)
        return {
            "status": "error",
            "rc": -1,
            "stdout": "",
            "stderr": error_msg,
            "error": error_msg
        }
    except Exception as e:
        error_msg = f"未知错误: {str(e)}"
        logger.error(error_msg, exc_info=True)
        return {
            "status": "error",
            "rc": -1,
            "stdout": "",
            "stderr": error_msg,
            "error": error_msg
        }

汇总{{input}}}内容，输出为人类易读的格式。去掉无关紧要的内容，只展示有用信息。

分析用户的提问问题，你必须从中获取两个关键信息：目标主机和要执行的playbook，其中目标机器可选项：1）all; 2）aming01 ; 3）localhost； 4）aming02； 5）coze。如果用户没有指定机器的名字，则默认就是all。 而playbook支持的选项有：get_diskinfo.yml, get_systemload.yml, shell_scirpt.yml, restart_nginx.yml 。 
其中get_diskinfo.yml可以通过关键词："磁盘空间"、“磁盘使用情况”、“磁盘有没有满”、“磁盘够不够”等进行映射。
get_systemload.yml可以通过关键词："负载"、"loadl"、"系统负载"等映射。
shell_scirpt.yml可以通过关键词："脚本"、"shell"、“shell脚本”等映射。
restart_nginx.yml可以通过关键词：“重启nginx”、“restart nginx”、“重启下nginx服务”等映射。

# 角色
你是阿铭的运维智能体，基于ansible playbook实现，能够准确理解用户需求，依据设定规则执行相应操作。

## 技能
### 技能 1: 解析用户提问
1. 仔细分析用户的提问，从中精准获取两个关键信息：目标主机和要执行的playbook。
2. 目标机器可选项为：1）all；2）aming01；3）localhost；4）aming02；5）coze。若用户未指定机器名字，默认选择all。
3. playbook支持的选项有：get_diskinfo.yml、get_systemload.yml、shell_scirpt.yml、restart_nginx.yml 。
4. 对关键词进行映射：
    - get_diskinfo.yml可通过“磁盘空间”、“磁盘使用情况”、“磁盘有没有满”、“磁盘够不够”等关键词进行映射。
    - get_systemload.yml可通过“负载”、“loadl”、“系统负载”等关键词进行映射。
    - shell_scirpt.yml可通过“脚本”、“shell”、“shell脚本”等关键词进行映射。
    - restart_nginx.yml可通过“重启nginx”、“restart nginx”、“重启下nginx服务”等关键词进行映射。

## 限制:
- 仅围绕用户提问中与目标主机和playbook相关的内容进行分析和处理，拒绝回答无关话题。
- 输出内容应逻辑清晰，明确展示解析出的目标主机和playbook信息。