faggot-chan/setup_wizard.py

#!/usr/bin/env python3
"""
setup wizard for faggot-chan. figures out which files need writing based on
how you're deploying, and writes them. run this before bot.py / main.py.
"""
import secrets
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent


def ask(prompt: str, default: str = None) -> str:
    suffix = f" [{default}]" if default else ""
    val = input(f"{prompt}{suffix}: ").strip()
    return val or default or ""


def ask_yn(prompt: str) -> bool:
    while True:
        val = input(f"{prompt} (y/n): ").strip().lower()
        if val in ("y", "yes"):
            return True
        if val in ("n", "no"):
            return False
        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_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_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://127.0.0.1:8765", ask_fallback: bool = False):
    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
    fallback = old[2] if len(old) > 2 else None
    print("\n-- connection to main.py --")
    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}")
    if ask_fallback:
        print("\n-- fallback (optional) --")
        if ask_yn("use a local ollama main.py as the fallback when gemini runs out of quota"):
            fallback = ask_keep("fallback main.py address", fallback or "ws://127.0.0.1:8766")
            print("start it with: PORT=8766 python3 main.py (it uses the same auth token, from config/auth_token.txt)")
        else:
            fallback = None
    (config_dir / "connection.txt").write_text(f"{uri}\n{token}\n" + (f"{fallback}\n" if fallback else ""), encoding="utf-8")
    print(f"wrote {config_dir / 'connection.txt'}")
    return token


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 --")
        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")

    colab = ask_yn("is this for a colab-hosted main.py")

    if not colab:
        # 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://127.0.0.1:8765", ask_fallback=(backend == "gemini"))
        write_auth_token_file(main_config, token=token)
        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")

    if where.lower().startswith("colab"):
        # 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)
        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),
    # this preset only wires up the connection to the colab-hosted main.py
    bot_config = SCRIPT_DIR / "files" / "config"
    if not (bot_config / "tokens.txt").exists():
        print("\nheads up: no files/config/tokens.txt found here — the vps preset skips")
        print("discord token setup, so make sure that file gets added separately.")
    write_connection_file(bot_config, default_uri="wss://your-colab-tunnel-url-here")
    print("\nall set. run bot.py once main.py is up on the colab side.")


if __name__ == "__main__":
    main()