import re, subprocess

env_path = "/etc/litellm/.env"

with open(env_path) as f:
    content = f.read()

# Add/replace key2
k2 = "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"
k3 = "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"

if "OPENROUTER_API_KEY_2=" in content:
    content = re.sub(r"OPENROUTER_API_KEY_2=.*", f"OPENROUTER_API_KEY_2={k2}", content)
else:
    content += f"\nOPENROUTER_API_KEY_2={k2}"

if "OPENROUTER_API_KEY_3=" in content:
    content = re.sub(r"OPENROUTER_API_KEY_3=.*", f"OPENROUTER_API_KEY_3={k3}", content)
else:
    content += f"\nOPENROUTER_API_KEY_3={k3}"

with open(env_path, "w") as f:
    f.write(content)

# Verify
with open(env_path) as f:
    for line in f:
        if "OPENROUTER" in line:
            parts = line.strip().split("=", 1)
            key = parts[0]
            val = parts[1] if len(parts) > 1 else ""
            print(f"{key}={val[:15]}...{val[-5:]}" if len(val) > 25 else f"{key}={val}")

subprocess.run(["systemctl", "restart", "ai-proxy"], capture_output=True)
print("\nProxy restarted!")