#!/usr/bin/env python3
"""打开 Dify 配置页上的「知识检索」。官方镜像把这个开关编译成关闭。

只改 dify-web-1 里的前端包，然后重启这一个容器。
不要重启 dify-api、New-API 或 vLLM。
容器如果被删除重建，镜像里的开关会回到关闭，需要再执行一次。
"""
import subprocess

OLD = 'ENABLE_AGENT_KNOWLEDGE_RETRIEVAL",0,!1'
NEW = 'ENABLE_AGENT_KNOWLEDGE_RETRIEVAL",0,!0'
ROOT = "/app/targets/next/web/.next"

listing = subprocess.check_output(
    ["docker", "exec", "dify-web-1", "grep", "-R", "-F", "-l", "ENABLE_AGENT_KNOWLEDGE_RETRIEVAL", ROOT],
    text=True,
)
changed = 0
for path in [line.strip() for line in listing.splitlines() if line.strip()]:
    local = "/tmp/" + path.strip("/").replace("/", "_")
    subprocess.check_call(["docker", "cp", f"dify-web-1:{path}", local])
    text = open(local, encoding="utf-8", errors="replace").read()
    count = text.count(OLD)
    if count == 0:
        print("skip", path, "already", text.count(NEW))
        continue
    open(local, "w", encoding="utf-8").write(text.replace(OLD, NEW))
    subprocess.check_call(["docker", "cp", local, f"dify-web-1:{path}"])
    print("patched", path, count)
    changed += count

if changed:
    subprocess.check_call(["docker", "restart", "dify-web-1"])
    print("restarted dify-web-1")
else:
    print("flag already enabled")