"""一键配置 BaoDan API Key 到 docker-compose.baodan.yml。

使用方式：
    python scripts/setup_bodan_api_keys.py

会提示输入三个 API Key，自动更新 docker-compose 配置。
"""
import re
import sys

COMPOSE_FILE = "docker-compose.baodan.yml"


def update_env(compose_path: str, key: str, value: str):
    """更新 docker-compose 中的环境变量。"""
    with open(compose_path, "r", encoding="utf-8") as f:
        content = f.read()

    # 匹配 KEY: "value" 或 KEY: "placeholder-..."
    pattern = rf'({re.escape(key)}:\s*")[^"]*(")'
    replacement = rf'\g<1>{value}\2'

    new_content, count = re.subn(pattern, replacement, content, count=1)
    if count == 0:
        print(f"  [WARN] 未找到 {key}，跳过")
        return False

    with open(compose_path, "w", encoding="utf-8") as f:
        f.write(new_content)
    return True


def main():
    print("=" * 50)
    print("  BaoDan API Key 配置工具")
    print("=" * 50)
    print()
    print("请在后台创建应用后，输入对应的 API Key。")
    print("如果暂时没有，可以直接回车跳过（用占位值）。")
    print()

    keys = {
        "BAODAN_CHAT_API_KEY": "对话应用 API Key（用于智能问答）",
        "BAODAN_WORKFLOW_API_KEY": "Workflow 应用 API Key（用于产品推荐）",
        "BAODAN_KB_API_KEY": "知识库 API Key（用于文档管理）",
    }

    compose_path = COMPOSE_FILE
    updated = 0

    for env_key, desc in keys.items():
        value = input(f"  {desc}\n  > {env_key}: ").strip()
        if value:
            if update_env(compose_path, env_key, value):
                print(f"  ✓ {env_key} 已更新")
                updated += 1
        else:
            print(f"  - 跳过 {env_key}")

    print()
    if updated > 0:
        print(f"完成！已更新 {updated} 个 API Key。")
        print("请运行以下命令重启服务：")
        print("  docker compose -f docker-compose.baodan.yml restart baodan-api baodan-worker")
    else:
        print("未更新任何配置。")


if __name__ == "__main__":
    main()