#!/usr/bin/env python3
"""
Bot Self-Modifier Script
通过 API 读取和修改 bot 配置：系统提示词、MCP服务器、技能、环境变量
"""

import argparse
import json
import os
import sys
import urllib.request
import urllib.error
import urllib.parse


def get_config():
    """获取配置，下面的MASTERKEY和ASSISTANT_ID是从环境变量自动获取的，不需要用户提供"""
    masterkey = os.environ.get("MASTERKEY", "master")
    bot_id = os.environ.get("ASSISTANT_ID", "")
    if not masterkey:
        print("ERROR: MASTERKEY environment variable is required")
        sys.exit(1)
    return masterkey, "http://localhost:8001", bot_id


def api_request(method, url, headers=None, data=None, is_multipart=False):
    """发送 HTTP 请求"""
    if headers is None:
        headers = {}

    if data is not None and not is_multipart:
        data = json.dumps(data).encode("utf-8")
        headers.setdefault("Content-Type", "application/json")

    req = urllib.request.Request(url, data=data, headers=headers, method=method)

    try:
        with urllib.request.urlopen(req) as resp:
            body = resp.read().decode("utf-8")
            return json.loads(body) if body else {}
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8")
        try:
            detail = json.loads(body)
        except Exception:
            detail = body
        print(f"ERROR: HTTP {e.code} - {detail}")
        sys.exit(1)


def build_multipart_data(fields, files):
    """构建 multipart/form-data 请求体"""
    boundary = "----BotModifierBoundary"
    lines = []

    for key, value in fields.items():
        lines.append(f"--{boundary}".encode())
        lines.append(f'Content-Disposition: form-data; name="{key}"'.encode())
        lines.append(b"")
        lines.append(value.encode() if isinstance(value, str) else value)

    for key, (filename, file_data, content_type) in files.items():
        lines.append(f"--{boundary}".encode())
        lines.append(
            f'Content-Disposition: form-data; name="{key}"; filename="{filename}"'.encode()
        )
        lines.append(f"Content-Type: {content_type}".encode())
        lines.append(b"")
        lines.append(file_data)

    lines.append(f"--{boundary}--".encode())
    lines.append(b"")

    body = b"\r\n".join(lines)
    content_type = f"multipart/form-data; boundary={boundary}"
    return body, content_type


# ============ 系统提示词 ============

def get_prompt(bot_id, masterkey, api_base):
    """读取系统提示词"""
    url = f"{api_base}/api/v1/bots/{bot_id}/settings"
    headers = {"Authorization": f"Bearer {masterkey}"}
    result = api_request("GET", url, headers)

    prompt = result.get("system_prompt", "")
    print("=== System Prompt ===")
    print(prompt if prompt else "(empty)")
    return prompt


def set_prompt(bot_id, masterkey, api_base, value):
    """修改系统提示词"""
    url = f"{api_base}/api/v1/bots/{bot_id}/settings"
    headers = {"Authorization": f"Bearer {masterkey}"}
    data = {"system_prompt": value}
    result = api_request("PUT", url, headers, data)
    print(f"OK: System prompt updated. {result.get('message', '')}")


# ============ Bot 基本信息 ============

def get_info(bot_id, masterkey, api_base):
    """读取 bot 基本信息（头像、描述、建议问题）"""
    url = f"{api_base}/api/v1/bots/{bot_id}/settings"
    headers = {"Authorization": f"Bearer {masterkey}"}
    result = api_request("GET", url, headers)

    print("=== Bot Info ===")
    print(f"  Name: {result.get('name', '(empty)')}")
    print(f"  Avatar URL: {result.get('avatar_url') or '(empty)'}")
    print(f"  Description: {result.get('description') or '(empty)'}")
    suggestions = result.get("suggestions") or []
    print(f"  Suggestions: {json.dumps(suggestions, ensure_ascii=False) if suggestions else '(empty)'}")
    return result


def set_avatar(bot_id, masterkey, api_base, value):
    """设置头像 URL"""
    url = f"{api_base}/api/v1/bots/{bot_id}/settings"
    headers = {"Authorization": f"Bearer {masterkey}"}
    data = {"avatar_url": value}
    result = api_request("PUT", url, headers, data)
    print(f"OK: Avatar updated. {result.get('message', '')}")


def set_name(bot_id, masterkey, api_base, value):
    """设置 bot 标题名称"""
    url = f"{api_base}/api/v1/bots/{bot_id}"
    headers = {"Authorization": f"Bearer {masterkey}"}
    data = {"name": value}
    result = api_request("PUT", url, headers, data)
    print(f"OK: Name updated to '{result.get('name', value)}'.")


def set_description(bot_id, masterkey, api_base, value):
    """设置描述"""
    url = f"{api_base}/api/v1/bots/{bot_id}/settings"
    headers = {"Authorization": f"Bearer {masterkey}"}
    data = {"description": value}
    result = api_request("PUT", url, headers, data)
    print(f"OK: Description updated. {result.get('message', '')}")


def set_suggestions(bot_id, masterkey, api_base, value):
    """设置建议问题（JSON 数组）"""
    try:
        suggestions = json.loads(value)
    except json.JSONDecodeError:
        print("ERROR: --value must be a valid JSON array for set_suggestions")
        sys.exit(1)

    if not isinstance(suggestions, list):
        print("ERROR: --value must be a JSON array, e.g. '[\"question1\", \"question2\"]'")
        sys.exit(1)

    url = f"{api_base}/api/v1/bots/{bot_id}/settings"
    headers = {"Authorization": f"Bearer {masterkey}"}
    data = {"suggestions": suggestions}
    result = api_request("PUT", url, headers, data)
    print(f"OK: Suggestions updated. {result.get('message', '')}")


# ============ MCP 服务器 ============

def list_mcp(bot_id, masterkey, api_base):
    """列出 MCP 服务器"""
    url = f"{api_base}/api/v1/bots/{bot_id}/mcp"
    headers = {"Authorization": f"Bearer {masterkey}"}
    result = api_request("GET", url, headers)

    if not result:
        print("No MCP servers configured.")
        return

    print("=== MCP Servers ===")
    for server in result:
        enabled_str = "enabled" if server.get("enabled") else "disabled"
        print(f"  [{enabled_str}] {server['name']} (id: {server['id']}, type: {server.get('type', 'N/A')})")
        config = server.get("config", {})
        if config:
            print(f"    config: {json.dumps(config, ensure_ascii=False)}")
    return result


def add_mcp(bot_id, masterkey, api_base, name, mcp_type, config_str):
    """添加 MCP 服务器"""
    try:
        config = json.loads(config_str)
    except json.JSONDecodeError:
        print("ERROR: --config must be valid JSON")
        sys.exit(1)

    url = f"{api_base}/api/v1/bots/{bot_id}/mcp"
    headers = {"Authorization": f"Bearer {masterkey}"}
    data = {
        "name": name,
        "type": mcp_type,
        "config": config,
        "enabled": True,
    }
    result = api_request("POST", url, headers, data)
    print(f"OK: MCP server '{name}' added (id: {result.get('id', 'N/A')})")


def delete_mcp(bot_id, masterkey, api_base, mcp_id):
    """删除 MCP 服务器"""
    url = f"{api_base}/api/v1/bots/{bot_id}/mcp/{mcp_id}"
    headers = {"Authorization": f"Bearer {masterkey}"}
    result = api_request("DELETE", url, headers)
    print(f"OK: MCP server deleted. {result.get('message', '')}")


# ============ 技能操作 ============

def list_skills(bot_id, masterkey, api_base):
    """列出所有技能"""
    # 获取技能文件列表
    url = f"{api_base}/api/v1/skill/list?bot_id={urllib.parse.quote(bot_id)}"
    headers = {"Authorization": f"Bearer {masterkey}"}
    result = api_request("GET", url, headers)

    # 获取当前启用的技能
    settings_url = f"{api_base}/api/v1/bots/{bot_id}/settings"
    settings = api_request("GET", settings_url, headers)
    enabled_skills_str = settings.get("skills", "") or ""
    enabled_skills = [s.strip() for s in enabled_skills_str.split(",") if s.strip()]

    skills = result.get("skills", [])
    print(f"=== Skills (total: {result.get('total', len(skills))}) ===")
    print(f"Enabled: {', '.join(enabled_skills) if enabled_skills else '(none)'}")
    print()
    for skill in skills:
        is_user = "[user]" if skill.get("user_skill") else "[official]"
        is_enabled = "*" if skill["name"] in enabled_skills else " "
        print(f"  {is_enabled} {is_user} {skill['name']}: {skill.get('description', '')}")
    return skills


def upload_skill(bot_id, masterkey, api_base, file_path):
    """上传技能 zip"""
    if not os.path.exists(file_path):
        print(f"ERROR: File not found: {file_path}")
        sys.exit(1)

    filename = os.path.basename(file_path)
    with open(file_path, "rb") as f:
        file_data = f.read()

    body, content_type = build_multipart_data(
        fields={"bot_id": bot_id},
        files={"file": (filename, file_data, "application/zip")},
    )

    url = f"{api_base}/api/v1/skill/upload"
    headers = {
        "Authorization": f"Bearer {masterkey}",
        "Content-Type": content_type,
    }
    result = api_request("POST", url, headers, body, is_multipart=True)
    print(f"OK: Skill uploaded. {result.get('message', '')} (name: {result.get('skill_name', 'N/A')})")


def enable_skill(bot_id, masterkey, api_base, skill_names):
    """启用技能 - 将技能名添加到 settings.skills"""
    # 先获取当前启用的技能
    settings_url = f"{api_base}/api/v1/bots/{bot_id}/settings"
    headers = {"Authorization": f"Bearer {masterkey}"}
    settings = api_request("GET", settings_url, headers)

    current = settings.get("skills", "") or ""
    current_list = [s.strip() for s in current.split(",") if s.strip()]

    new_skills = [s.strip() for s in skill_names.split(",") if s.strip()]
    for s in new_skills:
        if s not in current_list:
            current_list.append(s)

    updated = ",".join(current_list)
    data = {"skills": updated}
    result = api_request("PUT", settings_url, headers, data)
    print(f"OK: Skills enabled. Current: {updated}")


def disable_skill(bot_id, masterkey, api_base, skill_name):
    """禁用技能 - 从 settings.skills 中移除"""
    settings_url = f"{api_base}/api/v1/bots/{bot_id}/settings"
    headers = {"Authorization": f"Bearer {masterkey}"}
    settings = api_request("GET", settings_url, headers)

    current = settings.get("skills", "") or ""
    current_list = [s.strip() for s in current.split(",") if s.strip()]

    to_remove = [s.strip() for s in skill_name.split(",") if s.strip()]
    current_list = [s for s in current_list if s not in to_remove]

    updated = ",".join(current_list)
    data = {"skills": updated}
    result = api_request("PUT", settings_url, headers, data)
    print(f"OK: Skills disabled. Current: {updated if updated else '(none)'}")


def delete_skill(bot_id, masterkey, api_base, skill_name):
    """删除已上传的技能"""
    url = (
        f"{api_base}/api/v1/skill/remove"
        f"?bot_id={urllib.parse.quote(bot_id)}"
        f"&skill_name={urllib.parse.quote(skill_name)}"
    )
    headers = {"Authorization": f"Bearer {masterkey}"}
    result = api_request("DELETE", url, headers)
    print(f"OK: Skill '{skill_name}' deleted. {result.get('message', '')}")

    # 同时从启用列表中移除
    disable_skill(bot_id, masterkey, api_base, skill_name)


# ============ 环境变量 ============

def get_env(bot_id, masterkey, api_base):
    """读取环境变量"""
    url = f"{api_base}/api/v1/bots/{bot_id}/settings"
    headers = {"Authorization": f"Bearer {masterkey}"}
    result = api_request("GET", url, headers)

    shell_env = result.get("shell_env", {})
    print("=== Environment Variables ===")
    if not shell_env:
        print("(none)")
    else:
        for key, value in shell_env.items():
            display_val = value if value else "(empty)"
            print(f"  {key}={display_val}")
    return shell_env


def set_env(bot_id, masterkey, api_base, env_json):
    """修改环境变量（合并更新，不会覆盖未指定的变量）"""
    # 先获取当前环境变量
    settings_url = f"{api_base}/api/v1/bots/{bot_id}/settings"
    headers = {"Authorization": f"Bearer {masterkey}"}
    settings = api_request("GET", settings_url, headers)

    current_env = settings.get("shell_env", {}) or {}

    try:
        new_env = json.loads(env_json)
    except json.JSONDecodeError:
        print("ERROR: --value must be valid JSON for set_env")
        sys.exit(1)

    # 合并
    current_env.update(new_env)

    data = {"shell_env": current_env}
    result = api_request("PUT", settings_url, headers, data)
    print(f"OK: Environment variables updated. {result.get('message', '')}")
    for key, value in current_env.items():
        print(f"  {key}={value if value else '(empty)'}")


# ============ Main ============

def main():
    parser = argparse.ArgumentParser(description="Bot Self-Modifier: read and modify bot configuration via API")
    parser.add_argument(
        "--action",
        required=True,
        choices=[
            "get_prompt", "set_prompt",
            "get_info", "set_name", "set_avatar", "set_description", "set_suggestions",
            "list_mcp", "add_mcp", "delete_mcp",
            "list_skills", "upload_skill", "enable_skill", "disable_skill", "delete_skill",
            "get_env", "set_env",
        ],
        help="Action to perform",
    )
    parser.add_argument("--value", help="Value for the action (prompt text, skill name, env JSON, etc.)")
    parser.add_argument("--name", help="MCP server name (for add_mcp)")
    parser.add_argument("--mcp-type", default="sse", help="MCP server type (for add_mcp, default: sse)")
    parser.add_argument("--config", help="MCP server config JSON (for add_mcp)")
    parser.add_argument("--mcp-id", help="MCP server ID (for delete_mcp)")
    parser.add_argument("--file", help="Skill zip file path (for upload_skill)")

    args = parser.parse_args()
    masterkey, api_base, bot_id = get_config()

    action_map = {
        "get_prompt": lambda: get_prompt(bot_id, masterkey, api_base),
        "set_prompt": lambda: set_prompt(bot_id, masterkey, api_base, args.value),
        "get_info": lambda: get_info(bot_id, masterkey, api_base),
        "set_name": lambda: set_name(bot_id, masterkey, api_base, args.value),
        "set_avatar": lambda: set_avatar(bot_id, masterkey, api_base, args.value),
        "set_description": lambda: set_description(bot_id, masterkey, api_base, args.value),
        "set_suggestions": lambda: set_suggestions(bot_id, masterkey, api_base, args.value),
        "list_mcp": lambda: list_mcp(bot_id, masterkey, api_base),
        "add_mcp": lambda: add_mcp(bot_id, masterkey, api_base, args.name, args.mcp_type, args.config),
        "delete_mcp": lambda: delete_mcp(bot_id, masterkey, api_base, args.mcp_id),
        "list_skills": lambda: list_skills(bot_id, masterkey, api_base),
        "upload_skill": lambda: upload_skill(bot_id, masterkey, api_base, args.file),
        "enable_skill": lambda: enable_skill(bot_id, masterkey, api_base, args.value),
        "disable_skill": lambda: disable_skill(bot_id, masterkey, api_base, args.value),
        "delete_skill": lambda: delete_skill(bot_id, masterkey, api_base, args.value),
        "get_env": lambda: get_env(bot_id, masterkey, api_base),
        "set_env": lambda: set_env(bot_id, masterkey, api_base, args.value),
    }

    # 参数验证
    if args.action == "set_prompt" and not args.value:
        print("ERROR: --value is required for set_prompt")
        sys.exit(1)
    if args.action in ("set_avatar", "set_name", "set_description", "set_suggestions") and not args.value:
        print(f"ERROR: --value is required for {args.action}")
        sys.exit(1)
    if args.action == "add_mcp" and (not args.name or not args.config):
        print("ERROR: --name and --config are required for add_mcp")
        sys.exit(1)
    if args.action == "delete_mcp" and not args.mcp_id:
        print("ERROR: --mcp-id is required for delete_mcp")
        sys.exit(1)
    if args.action == "upload_skill" and not args.file:
        print("ERROR: --file is required for upload_skill")
        sys.exit(1)
    if args.action in ("enable_skill", "disable_skill", "delete_skill") and not args.value:
        print(f"ERROR: --value is required for {args.action}")
        sys.exit(1)
    if args.action == "set_env" and not args.value:
        print("ERROR: --value is required for set_env (JSON format)")
        sys.exit(1)

    action_map[args.action]()


if __name__ == "__main__":
    main()