bunch of stuff

1 files changed, 74 insertions, 18 deletions

diff --git a/setup_wizard.py b/setup_wizard.py
index 8575b27..743db51 100644
--- a/setup_wizard.py
+++ b/setup_wizard.py
@@ -25,27 +25,59 @@ def ask_yn(prompt: str) -> bool:
         print("just y or n please")
 
 
+def read_lines(path: Path) -> list[str]:
+    # same parsing bot.py does, empty list if the file isn't there yet
+    if not path.exists():
+        return []
+    return [l.strip() for l in path.read_text(encoding="utf-8").splitlines() if l.strip() and not l.strip().startswith("#")]
+
+
+def ask_keep(prompt: str, current: str = None, secret: bool = False) -> str:
+    # enter on an empty line keeps whatever the file already had (secrets don't get printed)
+    if current:
+        shown = "keep current" if secret else f"keep: {current}"
+        return input(f"{prompt} [{shown}]: ").strip() or current
+    return input(f"{prompt}: ").strip()
+
+
+def ask_backend() -> str:
+    while True:
+        val = ask("is main.py using 'ollama' or the 'gemini' api", "ollama").lower()
+        if val in ("ollama", "gemini"):
+            return val
+        print("just ollama or gemini please")
+
+
 def write_tokens_file(config_dir: Path):
     config_dir.mkdir(parents=True, exist_ok=True)
+    old = read_lines(config_dir / "tokens.txt")
+    old_token = old[0] if len(old) > 0 else None
+    old_server = old[1] if len(old) > 1 else None
     print("\n-- discord bot setup --")
-    token = ask("discord bot token")
-    server_id = ask("discord server (guild) id")
+    token = ask_keep("discord bot token", old_token, secret=True)
+    server_id = ask_keep("discord server (guild) id", old_server)
     while not server_id.isdigit():
         print("that's not a number")
-        server_id = ask("discord server (guild) id")
+        server_id = ask_keep("discord server (guild) id", old_server)
     (config_dir / "tokens.txt").write_text(f"{token}\n{server_id}\n", encoding="utf-8")
     print(f"wrote {config_dir / 'tokens.txt'}")
 
 
-def write_connection_file(config_dir: Path, default_uri: str = "ws://0.0.0.0:8765"):
+def write_connection_file(config_dir: Path, default_uri: str = "ws://127.0.0.1:8765"):
     config_dir.mkdir(parents=True, exist_ok=True)
+    old = read_lines(config_dir / "connection.txt")
+    old_uri = old[0] if len(old) > 0 else default_uri
+    old_token = old[1] if len(old) > 1 else None
     print("\n-- connection to main.py --")
-    uri = ask("main.py address (ws:// or wss://)", default_uri)
-    print("paste the auth token main.py is using (from colab_setup.py output),")
-    token = ask("or leave blank to generate a new one")
-    if not token:
-        token = secrets.token_urlsafe(32)
-        print(f"generated token: {token}")
+    uri = ask_keep("main.py address (ws:// or wss://)", old_uri)
+    if old_token:
+        token = ask_keep("auth token main.py is using (from colab_setup.py output)", old_token, secret=True)
+    else:
+        print("paste the auth token main.py is using (from colab_setup.py output),")
+        token = ask("or leave blank to generate a new one")
+        if not token:
+            token = secrets.token_urlsafe(32)
+            print(f"generated token: {token}")
     (config_dir / "connection.txt").write_text(f"{uri}\n{token}\n", encoding="utf-8")
     print(f"wrote {config_dir / 'connection.txt'}")
     return token
@@ -54,17 +86,33 @@ def write_connection_file(config_dir: Path, default_uri: str = "ws://0.0.0.0:876
 def write_auth_token_file(config_dir: Path, token: str = None):
     config_dir.mkdir(parents=True, exist_ok=True)
     if not token:
+        old = read_lines(config_dir / "auth_token.txt")
         print("\n-- websocket auth for main.py --")
-        print("this is the secret bot.py has to send. leave blank to generate one.")
-        token = ask("auth token")
-        if not token:
-            token = secrets.token_urlsafe(32)
-            print(f"generated token: {token}")
+        if old:
+            token = ask_keep("auth token", old[0], secret=True)
+        else:
+            print("this is the secret bot.py has to send. leave blank to generate one.")
+            token = ask("auth token")
+            if not token:
+                token = secrets.token_urlsafe(32)
+                print(f"generated token: {token}")
     (config_dir / "auth_token.txt").write_text(f"{token}\n", encoding="utf-8")
     print(f"wrote {config_dir / 'auth_token.txt'}")
     print(f"\ngive this token (and this machine's public wss:// url) to whoever's setting up bot.py:\n  {token}")
 
 
+def write_gemini_key_file(config_dir: Path):
+    config_dir.mkdir(parents=True, exist_ok=True)
+    old = read_lines(config_dir / "gemini_key.txt")
+    print("\n-- gemini api key for gemini.py --")
+    key = ask_keep("gemini api key", old[0] if old else None, secret=True)
+    while not key:
+        print("need a key")
+        key = ask_keep("gemini api key")
+    (config_dir / "gemini_key.txt").write_text(f"{key}\n", encoding="utf-8")
+    print(f"wrote {config_dir / 'gemini_key.txt'}")
+
+
 def main():
     print("faggot-chan setup wizard\n")
 
@@ -74,10 +122,13 @@ def main():
         # everything lives together, plain local setup
         bot_config = SCRIPT_DIR / "files" / "config"
         main_config = SCRIPT_DIR / "config"
+        backend = ask_backend()
         write_tokens_file(bot_config)
-        token = write_connection_file(bot_config, default_uri="ws://0.0.0.0:8765")
+        token = write_connection_file(bot_config, default_uri="ws://127.0.0.1:8765")
         write_auth_token_file(main_config, token=token)
-        print("\nall set. run main.py, then bot.py.")
+        if backend == "gemini":
+            write_gemini_key_file(main_config)
+        print(f"\nall set. run {'gemini.py' if backend == 'gemini' else 'main.py'}, then bot.py.")
         return
 
     where = ask("running this on the 'vps' (bot.py host) or 'colab' box (main.py host)", "vps")
@@ -86,8 +137,13 @@ def main():
         # this box just needs the auth token main.py will check — connection details
         # (discord token, server id) live on whichever box runs bot.py
         main_config = SCRIPT_DIR / "config"
+        backend = ask_backend()
         write_auth_token_file(main_config)
-        print("\nnow run colab_setup.py to bring main.py + the tunnel up.")
+        if backend == "gemini":
+            write_gemini_key_file(main_config)
+            print("\nnow run gemini.py (colab_setup.py would start the ollama main.py and pull models you don't need).")
+        else:
+            print("\nnow run colab_setup.py to bring main.py + the tunnel up.")
         return
 
     # vps / bot.py host: skip discord-client setup entirely (handled separately),