#!/usr/bin/env python3
"""An InfiTalk bot that connects through the gateway wss://tt.infiartt.com/bot.

Python 3.9 or newer with one package:  pip install websockets

  1. Make a bot on https://infiartt.com/dev-center/me and copy its token (itw_...).
  2. Add it to a server you manage ("Add to a server" on the same page).
  3. INFITALK_TOKEN=itw_... python echo_bot.py      (PowerShell: $env:INFITALK_TOKEN="itw_..."; python echo_bot.py)

The same flow as echo-bot.js: identify, heartbeats, resume, a menu, a dialog, a reply when named and a
quiet one-second tone. Docs: https://infiartt.com/dev-center/infitalk-bots
"""
import asyncio
import base64
import json
import math
import os
import struct
import sys

import websockets

GATEWAY = os.environ.get("INFITALK_GATEWAY", "wss://tt.infiartt.com/bot")
TOKEN = os.environ.get("INFITALK_TOKEN", "")
CHANNEL = os.environ.get("INFITALK_CHANNEL", "/")  # the channel it starts in, by path
if not TOKEN.startswith("itw_"):
    sys.exit("Set INFITALK_TOKEN to the token of your bot (itw_...).")

MENU = [
    {"name": "echo", "description": "Repeat what you write", "category": "Fun",
     "options": [{"name": "text", "label": "Text", "type": "text", "required": True}]},
    {"name": "hello", "description": "Open a small dialog", "category": "Fun", "dialog": True},
    {"name": "tone", "description": "Play a quiet tone in your channel for one second", "category": "Sound"},
]

resume = {}  # server -> {"session", "seq"}: nothing is missed after a short break


async def send(ws, msg):
    await ws.send(json.dumps(msg))


async def heartbeat(ws, every_ms):
    while True:
        await asyncio.sleep(every_ms / 1000)
        await send(ws, {"op": "heartbeat"})


async def play_tone(ws, server, channel):
    """48 kHz stereo 16-bit PCM in 20 ms pieces, in real time. Quiet on purpose (about -32 dB)."""
    rate, piece, amp, phase = 48000, 960, 800, 0.0
    for _ in range(50):
        frames = bytearray()
        for _ in range(piece):
            v = int(math.sin(phase) * amp)
            phase += 2 * math.pi * 440 / rate
            frames += struct.pack("<hh", v, v)
        await send(ws, {"server": server, "op": "audio", "channel": channel, "pcm": base64.b64encode(frames).decode()})
        await asyncio.sleep(0.02)
    await send(ws, {"server": server, "op": "audio_end", "channel": channel})


async def on_command(ws, m):
    server, interaction = m["server"], m["interaction"]
    if m["command"] == "echo":
        await send(ws, {"server": server, "op": "respond", "interaction": interaction,
                        "content": m["values"]["text"], "public": True})
    elif m["command"] == "hello":
        await send(ws, {"server": server, "op": "dialog", "interaction": interaction, "dialog": {
            "id": "hello", "title": "Hello", "focus": "name",
            "controls": [
                {"type": "text", "id": "name", "label": "&Your name", "value": m["user"]["nickname"], "maxlength": 40},
                {"type": "button", "id": "greet", "label": "&Greet me", "default": True},
                {"type": "button", "id": "cancel", "label": "Cancel", "close": True},
            ]}})
    elif m["command"] == "tone":
        await send(ws, {"server": server, "op": "respond", "interaction": interaction,
                        "content": "Playing a quiet tone for one second."})
        await send(ws, {"server": server, "op": "join", "channel": m["channel"]})  # audio needs the bot there
        asyncio.create_task(play_tone(ws, server, m["channel"]))


async def run():
    async with websockets.connect(GATEWAY, max_size=2 ** 20) as ws:
        beat = None
        try:
            async for text in ws:
                m = json.loads(text)
                server = m.get("server")
                if server is not None and isinstance(m.get("seq"), int) and str(server) in resume:
                    resume[str(server)]["seq"] = m["seq"]
                op = m.get("op")
                if op == "hello":
                    beat = asyncio.create_task(heartbeat(ws, m["heartbeat_interval"]))
                    await send(ws, {"op": "identify", "token": TOKEN, "resume": resume})
                elif op == "identified":
                    names = ", ".join(f"{s['name']} ({'awake' if s['awake'] else 'asleep'})" for s in m["servers"])
                    print(f"Signed in as {m['bot']['name']}; servers: {names or 'none yet'}")
                elif op == "ready":
                    old = resume.get(str(server), {}).get("seq", 0)
                    resume[str(server)] = {"session": m["session"], "seq": old if m.get("resumed") else 0}
                    print(f"Server {server} ready ({'resumed' if m.get('resumed') else 'new session'})")
                    await send(ws, {"server": server, "op": "commands", "commands": MENU})
                    # people use a bot's menu in the channel it is in: start in the main channel
                    if not m.get("resumed") and not m["bot"].get("channel"):
                        await send(ws, {"server": server, "op": "join", "path": CHANNEL})
                elif op == "server_sleeping":
                    print(f"Server {server} went to sleep; the gateway connects again when it wakes")
                elif op == "command":
                    await on_command(ws, m)
                elif op == "dialog_event" and m.get("event") == "button" and m.get("control") == "greet":
                    name = (m.get("values") or {}).get("name") or m["user"]["nickname"]
                    await send(ws, {"server": server, "op": "respond", "interaction": m["interaction"], "content": f"Hello, {name}!"})
                    await send(ws, {"server": server, "op": "dialog_close", "instance": m["instance"]})
                elif op == "message" and m.get("mention"):
                    await send(ws, {"server": server, "op": "say", "channel": m["channel"], "reply_to": m["msgid"],
                                    "content": f"Hi {m['user']['nickname']}, try my menu with Ctrl+Shift+P."})
                elif op == "result" and not m.get("ok"):
                    print(f"Refused ({server if server is not None else 'gateway'}): {m.get('error')}")
        finally:
            if beat:
                beat.cancel()


async def main():
    while True:
        try:
            await run()
        except websockets.ConnectionClosed as e:
            code = e.rcvd.code if e.rcvd else None
            if code == 4004:
                sys.exit("The token is not valid any more, or the bot was turned off.")
            print(f"Disconnected ({code}); connecting again in 5 seconds")
        except OSError as e:
            print(f"Cannot connect ({e}); trying again in 5 seconds")
        await asyncio.sleep(5)


if __name__ == "__main__":
    asyncio.run(main())