#!/usr/bin/env python3
"""Run from Dify 1.17.1 docker/. Existing secrets are retained on URL changes."""
import base64
import os
import re
import secrets
import sys
from pathlib import Path
from urllib.parse import urlsplit

if len(sys.argv) != 2:
    raise SystemExit("Usage: python3 prepare-env.py http://localhost:8080 OR https://dify.example.com")
base = sys.argv[1].rstrip("/")
url = urlsplit(base)
if url.path or url.query or url.fragment or url.username or not url.hostname:
    raise SystemExit("Use an origin without a path, credentials, query or fragment")
if url.scheme != "https" and base != "http://localhost:8080":
    raise SystemExit("HTTPS is required except for the initial SSH tunnel")
env_path = Path(".env")
fresh = not env_path.exists()
content = (Path(".env.example") if fresh else env_path).read_text()
values = {}
if fresh:
    for key in ("SECRET_KEY", "DB_PASSWORD", "REDIS_PASSWORD", "PLUGIN_DAEMON_KEY",
                "PLUGIN_DIFY_INNER_API_KEY", "SANDBOX_API_KEY", "DIFY_AGENT_API_TOKEN",
                "DIFY_AGENT_SERVER_SECRET_KEY", "DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN"):
        values[key] = secrets.token_hex(32)
    values["CELERY_BROKER_URL"] = "redis://:" + values["REDIS_PASSWORD"] + "@redis:6379/1"
    values["DIFY_AGENT_SERVER_SECRET_KEY"] = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode().rstrip("=")
for key in ("CONSOLE_API_URL", "CONSOLE_WEB_URL", "SERVICE_API_URL", "APP_API_URL",
            "APP_WEB_URL", "FILES_URL", "TRIGGER_URL"):
    values[key] = base
values["ENDPOINT_URL_TEMPLATE"] = base + "/e/{hook_id}"
values["NEXT_PUBLIC_SOCKET_URL"] = ("wss" if url.scheme == "https" else "ws") + "://" + url.netloc
for key, value in values.items():
    pattern = re.compile(r"^" + re.escape(key) + r"=.*$", re.MULTILINE)
    if pattern.search(content):
        content = pattern.sub(lambda match: key + "=" + value, content)
    else:
        content += "\n" + key + "=" + value + "\n"
env_path.write_text(content)
os.chmod(env_path, 0o600)
print("Created .env" if fresh else "Updated URLs; retained existing secrets")